{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wprowadzenie\n",
    "\n",
    "W tym notatniku przeanalizujemy splot GAT (ang. *Graph Attention Network*). Jest to warstwa splotu, której działanie bezpośrednio odpowiada mechanizmowi uwagi, stosowanemu w sieciach typu Transformer.\n",
    "\n",
    "Oryginalna publikacja:\n",
    "\n",
    "> Veličković, P., Cucurull, G., Casanova, A., Romero, A., Lio, P., & Bengio, Y. (2017). Graph attention networks. arXiv preprint arXiv:1710.10903."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import bibliotek"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "import torch as th\n",
    "import torch.nn as nn\n",
    "import torch_geometric as pyg\n",
    "import torch_geometric.nn as pyg_nn\n",
    "import torch_geometric.transforms as pyg_t\n",
    "import torch_geometric.datasets as pyg_datasets\n",
    "import torch_geometric.data as pyg_data\n",
    "\n",
    "from pathlib import Path\n",
    "from tqdm.auto import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import utilities as utils\n",
    "from base_models import NodesClassifierModel, NaiveGraphConv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Przygotowanie danych"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(x=[2708, 1433], edge_index=[2, 10556], y=[2708], train_mask=[2708], val_mask=[2708], test_mask=[2708])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = pyg_datasets.Planetoid(root=Path(os.pardir, \"../data\", \"Planetoid\"), name='Cora', transform=pyg_t.NormalizeFeatures())\n",
    "cora = dataset[0]\n",
    "cora"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GAT\n",
    "\n",
    "Warstwa splotu GAT odpowiada mechanizmowi uwagi, znanemu m.in. z sieci typu Transformer, stosowanych w analizie języka i modelach generatywnych.\n",
    "\n",
    "W przypadku splotu GAT, mechanizm działania jest następujący:\n",
    "1. Dla każdego wierzchołka v w grafie, wyszukiwane jest sąsiedztwo N(v).\n",
    "2. Dla każdego sąsiada u, wyliczany jest współczynnik uwagi $e_{vu}$ pomiędzy wierzchołkami v i u.\n",
    "3. Następnie, współczynniki uwagi są normalizowane, tak aby suma współczynników dla każdego wierzchołka była równa 1.\n",
    "4. Współczynniki uwagi są wykorzystywane do obliczenia wiadomości od wierzchołka u do wierzchołka v.\n",
    "\n",
    "Poniżej przedstawiona jest formalizacja splotu GAT.\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\mathbf{e}_{uv} & =\\text{LeakyReLU}\\left(\\text{Att}\\left([\\mathbf{W}^{(k)}\\mathbf{h}_{u}^{(k-1)}, \\mathbf{W}^{(k)}\\mathbf{h}_{v}^{(k-1)}] \\right)\\right) & \\triangleright \\text{Nieznormalizowane wsp. uwagi} \\\\[1em]\n",
    "\n",
    "\\mathbf{\\alpha}_{uv}&=\\text{Softmax}\\left(\\mathbf{e}_{uv}\\right) = \\frac{\\mathbf{e}_{uv}}{\\sum_{u\\in \\mathcal{N}(v)}\\mathbf{e}_{uv}} & \\triangleright \\text{Normalizacja wsp. uwagi} \\\\[1em]\n",
    "\n",
    "\\mathbf{m}_{u}^{(k)} &= \\bigoplus \\left\\{\\mathbf{\\alpha}_{uv}[h] \\mathbf{W}^{(k)}\\mathbf{h}_{u}^{(k-1)} \\right\\} & \\triangleright \\text{Tworzenie wiadomości ze wsp. uwagi} \\\\[1em]\n",
    "\n",
    "\\mathbf{h}_{\\mathcal{N}(v)}^{(k)} &= \\text{ŚREDNIA}^{(k)} \\left\\{\\mathbf{m}_{u}^{(k)} \\mid \\forall u \\in \\mathcal{N}(v)\\right\\} & \\triangleright \\text{Agregacja sąsiedztwa} \\\\[1em]\n",
    "\n",
    "\\mathbf{h}_v^{(k)} &= \\sigma \\left(\\mathbf{h}_{\\mathcal{N}(v)}^{(k)} \\right) & \\triangleright \\text{Aktualizacja wierzchołka} \\\\[1em]\n",
    "\\end{align*}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementacja naiwna"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NaiveGAT(NaiveGraphConv):\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            dim_in: int,\n",
    "            dim_out: int,\n",
    "            heads: int,\n",
    "            sigma: nn.Module = nn.Identity(),\n",
    "            neg_slope: float = 0.2,\n",
    "            agg_fn=pyg_nn.MeanAggregation(),\n",
    "            add_self_loops: bool = True):\n",
    "        super().__init__(\n",
    "            dim_in, \n",
    "            dim_out, \n",
    "            aggr=agg_fn, \n",
    "            sigma=sigma, \n",
    "            add_self_loops=add_self_loops)\n",
    "        self.dim_in = dim_in\n",
    "        self.dim_out = dim_out\n",
    "        self.neg_slope = neg_slope\n",
    "        self.heads = heads\n",
    "        self.W = nn.Linear(dim_in, dim_out, bias=False)\n",
    "        self.bias = nn.Parameter(th.empty(1, dim_out * heads))\n",
    "        self.att = nn.Linear(2*dim_out, heads)\n",
    "        self.reset_parameters()\n",
    "\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        pyg_nn.inits.glorot(self.W.weight)\n",
    "        pyg_nn.inits.glorot(self.att.weight)\n",
    "        nn.init.zeros_(self.bias)\n",
    "\n",
    "    def message(self, v_neighbors: th.Tensor, alpha: th.Tensor) -> th.Tensor:\n",
    "        \"\"\"Funkcja budująca wiadomości od sąsiadów, implementacja abstrakcyjnej\n",
    "        funkcji `message` z architektury MPNN. Przyjmuje jako wejście\n",
    "        macierz cech sąsiadów wierzchołka v:  `u in N(v)` i buduje z niej wiadomości.\n",
    "        W tym przypadku - wiadomości od sąsiadów są wymnażane przez znormalizowany\n",
    "        współczynnik uwagi dla każdej głowy.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        v_neighbors : th.Tensor\n",
    "            Macierz cech sąsiadów wierzchołka v: `N(v)`.\n",
    "            Wymiar: `|N(v)| x dim_in`.\n",
    "\n",
    "        alpha: th.Tensor\n",
    "            Współczynniki uwagi dla każdego sąsiada.\n",
    "            Wymiar: `|N(v)|`.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        th.Tensor\n",
    "            Macierz wiadomości od sąsiadów.\n",
    "            Wymiar: `|N(v)| x dim_out`.\n",
    "        \"\"\"\n",
    "        result = []\n",
    "        # Dla każdej głowy uwagi\n",
    "        for h in range(self.heads):\n",
    "            # Pomnóż reprezentację sąsiada przez współczynnik uwagi\n",
    "            # Wymiar: |N(v)| x dim_out\n",
    "            head_coefficients = v_neighbors * alpha[:, h].unsqueeze(-1)\n",
    "            result.append(head_coefficients)\n",
    "            \n",
    "        \"\"\"Powyższe można by zaimplementować w całości, w jednej linijce za pomocą wektoryzacji: \n",
    "        \n",
    "        result = v_neighbors.unsqueeze(1) * alpha.unsqueeze(-1)\n",
    "        \n",
    "        v_neighbors.unsqueeze(1) - rozszerzenie wymiaru o jedną osiąg, aby móc przeprowadzić operację mnożenia.\n",
    "                                   Nowy kształt: N x 1 x f\n",
    "        \n",
    "        alpha.unsqueeze(-1) - rozszerzenie wymiaru o jedną osiąg, aby móc przeprowadzić operację mnożenia.\n",
    "                                Nowy kształt: N x h x 1\n",
    "                                \n",
    "        (N x 1 x f) * (N x h x 1) = N x h x f\n",
    "        \n",
    "        Taka wersja jest jednak sporo mniej czytelna i zrozumiała, za to bardziej efektywna      \n",
    "        \"\"\"\n",
    "            \n",
    "        # Zestawienie wyników w jedną macierz, wymiar: |N(v)| x heads x dim_out\n",
    "        return th.stack(result, dim=-1).transpose(1, 2)\n",
    "\n",
    "    def aggregate(self, m_nv: th.Tensor) -> th.Tensor:\n",
    "        \"\"\"Funkcja agregująca wiadomości od sąsiadów, implementacja abstrakcyjnej\n",
    "        funkcji `aggregate` z architektury MPNN. Przyjmuje jako wejście\n",
    "        macierz wiadomości od sąsiadów wierzchołka v i agreguje je - wyliczając średnią wartość.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        m_nv : th.Tensor\n",
    "            Macierz wiadomości od sąsiadów.\n",
    "            Wymiar: `|N(v)| x dim_out`.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        th.Tensor\n",
    "            Zagregowane wiadomości. Wymiar: `1 x heads x dim_out`.\n",
    "        \"\"\"\n",
    "        # Uśrednienie wyników dla każdej głowy uwagi \n",
    "        return self.aggr(m_nv, dim=0)\n",
    "\n",
    "    def update(self, v: th.Tensor, agg: th.Tensor) -> th.Tensor:\n",
    "        \"\"\"Implementacja funkcji `update` z architektury MPNN. Aktualizuje\n",
    "        reprezentacje wierzchołka v na podstawie zagregowanych wiadomości\n",
    "        od sąsiadów. W przypadku splotu GAT - wykorzystuje uśrednione wartości\n",
    "        wiadomości sąsiadów, wliczając to pętlę wierzchołka v do samego siebie.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        v : th.Tensor\n",
    "            Aktualna reprezentacja wierzchołka v. Wymiar: `1 x dim_in`.\n",
    "        agg : th.Tensor\n",
    "            Zagregowane wiadomości sąsiadów. Wymiar: `1 x dim_out`.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        th.Tensor\n",
    "            Zaktualizowana reprezentacja wierzchołka v. Wymiar: `1 x dim_out`.\n",
    "        \"\"\"\n",
    "        # Aktualizacja reprezentacji wierzchołka za pomocą transformacji sigma\n",
    "        # Najpierw: \"spłaszczenie\" macierzy zagregowanych wiadomości\n",
    "        # Wymiar po spłaszczeniu: 1 x (heads * dim_out)\n",
    "        new_repr = self.sigma(agg.view(-1, self.heads * self.dim_out))\n",
    "        return new_repr\n",
    "    \n",
    "    def build_attention_coefs(self, X_v: th.Tensor, neighbors_X: th.Tensor) -> th.Tensor:\n",
    "        \"\"\"Funkcja odpowiedzialna za zbudowanie współczynników uwagi dla każdego\n",
    "        połączenia między wierzchołkiem v a jego sąsiadami.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        X_v : th.Tensor\n",
    "            Wektor reprezentacji wierzchołka v. Wymiar: 1 x dim_out\n",
    "        neighbors_X : th.Tensor\n",
    "            Macierz reprezentacji sąsiadów wierzchołka v. wymiar: |N(v)| x dim_out\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        th.Tensor\n",
    "            Współczynniki uwagi dla każdego połączenia między wierzchołkiem v a jego sąsiadami.\n",
    "            Wymiar: |N(v)| x heads\n",
    "        \"\"\"\n",
    "        # \"Rozciągnięcie\" albo powtórzenie wektora wierzchołka v\n",
    "        # tyle razy, ilu ma sąsiadów\n",
    "        X_v_expanded = X_v.expand(neighbors_X.shape[0], -1)\n",
    "            \n",
    "        # Konkatenacja - wymiar: |N(v)| x 2 * dim_out\n",
    "        X_uv = th.cat((X_v_expanded, neighbors_X), dim=-1)\n",
    "            \n",
    "        # Przetworzenie przez warstwę uwagi. Wymiar: |N(v)| x heads\n",
    "        attention_input = self.att(X_uv)\n",
    "        attention_coefs = nn.functional.leaky_relu(\n",
    "            attention_input, \n",
    "            negative_slope=self.neg_slope)\n",
    "            \n",
    "        # ZNORMALIZOWANE WSPÓŁCZYNNIKI UWAGI. Wymiar: |N(v)| x heads\n",
    "        alpha = nn.functional.softmax(attention_coefs, dim=0)\n",
    "        return alpha\n",
    "\n",
    "    def forward(self, x: th.Tensor, edge_index: th.Tensor) -> th.Tensor:\n",
    "        \"\"\"Funkcja przepływu danych w warstwie splotu grafowego. Rozpoczyna cykl\n",
    "        przekazywania wiadomości, dla każdego wierzchołka w grafie: v po kolei:\n",
    "        1. Znajdując sąsiadów;\n",
    "        2. Budując wiadomości od sąsiadów;\n",
    "        3. Agregując wiadomości od sąsiadów;\n",
    "        4. Aktualizując reprezentację wierzchołka.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        x : th.Tensor\n",
    "            Macierz cech wierzchołków. Wymiar: N x dim_in\n",
    "\n",
    "        edge_index : th.Tensor\n",
    "            Macierz sąsiedztwa grafu. Wymiar: 2 x E\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        th.Tensor\n",
    "            Zaktualizowane cechy wierzchołków.\n",
    "        \"\"\"\n",
    "        # Alokacja nowej macierzy reprezentacji wierzchołków\n",
    "        N = x.shape[0]\n",
    "        h_new = []\n",
    "\n",
    "        # Wymiar: N x dim_out\n",
    "        X_mapped = self.W(x)\n",
    "\n",
    "        # Dla każdego wierzchołka v w grafie\n",
    "        for v in range(x.shape[0]):\n",
    "            # Znajdź indeksy w macierzy sąsiedztwa, gdzie odbiera on wiadomość\n",
    "            v_index_as_receiver = th.where(edge_index[1] == v)[0]\n",
    "\n",
    "            # Znajdź sąsiadów, którzy wysyłają tę wiadomość\n",
    "            # Wymiar: |N(v)| x dim_in\n",
    "            neighbors = edge_index[0, v_index_as_receiver]\n",
    "\n",
    "            # Zbuduj wiadomości od sąsiadów\n",
    "            # Wymiar: |N(v)| x dim_out\n",
    "            neighbors_X = X_mapped[neighbors]\n",
    "\n",
    "            if self.add_self_loops:\n",
    "                neighbors_X = self.apply_self_loops(X_mapped, v, neighbors)\n",
    "\n",
    "            X_v = X_mapped[v].unsqueeze(0)\n",
    "            \n",
    "            alpha = self.build_attention_coefs(X_v, neighbors_X)\n",
    "\n",
    "            # Wiadomość od sąsiadó. Wymiar: |N(v)| x heads x dim_out\n",
    "            m_nv = self.message(neighbors_X, alpha=alpha)\n",
    "\n",
    "            # Agreguj wiadomości sąsiadów\n",
    "            # 1 x heads x dim_out\n",
    "            agg = self.aggregate(m_nv)\n",
    "\n",
    "            # Zaktualizuj wierzchołek v\n",
    "            # 1 x (heads*dim_out)\n",
    "            new_v = self.update(X_mapped[v], agg)\n",
    "            h_new.append(new_v)\n",
    "        h_new = th.cat(h_new) + self.bias\n",
    "        return h_new\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sprawdzenie, czy implementacja działa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2708, 32])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with th.no_grad():\n",
    "    naive_gat = NaiveGAT(dim_in=cora.num_features, dim_out=16, heads=2)\n",
    "    naive_gat.eval()\n",
    "    h = naive_gat(cora.x, cora.edge_index)\n",
    "h.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trywialny przykład\n",
    "\n",
    "Przykład zaprezentowany poniżej, na podstawie bardzo prostego grafu i ręcznie ustawianych wag powinien pomóc w zrozumieniu, w jaki sposób funkcjonuje warstwa splotu GAT."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bardzo prosty graf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 21852 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "Gdata = pyg_data.Data()\n",
    "Gdata.x = th.tensor(np.array([\n",
    "    [0.1, 0.2, 0.3],\n",
    "    [0.3, 0.4, 0.5],\n",
    "    [0.2, 0.2, 0.2]\n",
    "])).to(th.float)\n",
    "\n",
    "Gdata.edge_index = th.tensor([\n",
    "    [0, 1],\n",
    "    [2, 2]\n",
    "])\n",
    "\n",
    "Gnx = pyg.utils.to_networkx(Gdata)\n",
    "nx.draw(Gnx, with_labels=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Wszystkie wierzchołki łączą się z wierzchołkiem nr 2. Ustawimy warstwę splotu w sposób następujący:\n",
    "\n",
    "1. Ilość \"głów\" uwagi = 2.\n",
    "    1. Wagi pierwszej głowy ustawimy na 0.5\n",
    "    2. Wagi drugiej głowy ustawimy na 2.\n",
    "2. Ilość neuronów wyjściowych w warstwie projekcji = 3. Dla ułatwienia wszystkie wagi tej warstwy ustawimy na 1.\n",
    "3. Bias ustawimy na 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "with th.no_grad():\n",
    "    trivial_gat = NaiveGAT(dim_in=3, dim_out=3, heads=2)\n",
    "    trivial_gat.W.weight = nn.Parameter(th.ones(size=(3, 3)).to(th.float))\n",
    "    \n",
    "    trivial_gat.att.weight = nn.Parameter(th.tensor([\n",
    "        [0.5, 0.5, 0.5, 0.5, 0.5, 0.5], \n",
    "        [2., 2., 2., 2., 2., 2.]]).to(th.float))\n",
    "    trivial_gat.att.bias = nn.Parameter(th.zeros_like(trivial_gat.att.bias).to(th.float))\n",
    "    trivial_gat.bias = nn.Parameter(th.zeros(size=(1, 6)).to(th.float))\n",
    "    \n",
    "    out = trivial_gat(Gdata.x, Gdata.edge_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.6000, 0.6000, 0.6000, 0.6000, 0.6000, 0.6000],\n",
       "        [1.2000, 1.2000, 1.2000, 1.2000, 1.2000, 1.2000],\n",
       "        [0.3103, 0.3103, 0.3103, 0.3896, 0.3896, 0.3896]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Klasyfikacja wierzchołków"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model klasyfikacji z naiwną implementacją GAT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "⚠️ Uwaga, ze względu na fakt, iż implementacja naiwna nie jest zoptymalizowana, może być konieczne dłuższe oczekiwanie na wyniki. W zależności od posiadanego sprzętu i procesora, całość może zająć chwilę. Jeśli czas oczekiwani jest za długi - warto zmniejszyć ilość epok.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyg.seed_everything(123)\n",
    "\n",
    "nheads = 2\n",
    "out1 = 128\n",
    "out2 = 64\n",
    "\n",
    "\n",
    "naive_gat1 = NaiveGAT(\n",
    "    dim_in=cora.num_features,\n",
    "    dim_out=out1,\n",
    "    add_self_loops=True,\n",
    "    heads=nheads)\n",
    "\n",
    "naive_gat2 = NaiveGAT(\n",
    "    dim_in=out1 * nheads,\n",
    "    dim_out=out2,\n",
    "    add_self_loops=True,\n",
    "    heads=nheads)\n",
    "\n",
    "naive_clf = NodesClassifierModel(\n",
    "    gnn_convs=(naive_gat1, naive_gat2), \n",
    "    lin_sizes=(out2 * nheads, 64, 7))\n",
    "\n",
    "optim = th.optim.Adam(naive_clf.parameters(), lr=0.01)\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "with th.no_grad():\n",
    "    naive_clf.eval()\n",
    "    yhat_naive_before_training = naive_clf(cora.x, cora.edge_index)\n",
    "\n",
    "\n",
    "epoch_losses_naive = []\n",
    "val_losses_naive = []\n",
    "\n",
    "for epoch in tqdm(range(30)):\n",
    "    naive_clf.train()\n",
    "    optim.zero_grad()\n",
    "    out = naive_clf(cora.x, cora.edge_index)\n",
    "    if out.isnan().any():\n",
    "        raise ValueError(\"NaNs in the output.\")\n",
    "    loss = loss_fn(out[cora.train_mask], cora.y[cora.train_mask])\n",
    "    loss.backward()\n",
    "    optim.step()\n",
    "    epoch_losses_naive.append(loss.item())\n",
    "\n",
    "    naive_clf.eval()\n",
    "    with th.no_grad():\n",
    "        val_loss = loss_fn(out[cora.val_mask], cora.y[cora.val_mask])\n",
    "        val_losses_naive.append(val_loss.item())\n",
    "\n",
    "with th.no_grad():\n",
    "    naive_clf.eval()\n",
    "    yhat_naive_after_training = naive_clf(cora.x, cora.edge_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wynik klasyfikacji modelu Naiwny GAT przed trenowaniem:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.00      0.00      0.00        61\n",
      "           1       0.00      0.00      0.00        36\n",
      "           2       0.00      0.00      0.00        78\n",
      "           3       0.00      0.00      0.00       158\n",
      "           4       0.16      1.00      0.28        81\n",
      "           5       0.00      0.00      0.00        57\n",
      "           6       0.00      0.00      0.00        29\n",
      "\n",
      "    accuracy                           0.16       500\n",
      "   macro avg       0.02      0.14      0.04       500\n",
      "weighted avg       0.03      0.16      0.05       500\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "Wynik klasyfikacji modelu Naiwny GAT po trenowaniu:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.33      0.25      0.28        61\n",
      "           1       0.00      0.00      0.00        36\n",
      "           2       0.86      0.92      0.89        78\n",
      "           3       0.74      0.70      0.72       158\n",
      "           4       0.60      0.77      0.67        81\n",
      "           5       0.61      0.82      0.70        57\n",
      "           6       0.46      0.66      0.54        29\n",
      "\n",
      "    accuracy                           0.65       500\n",
      "   macro avg       0.51      0.59      0.54       500\n",
      "weighted avg       0.60      0.65      0.62       500\n",
      "\n"
     ]
    }
   ],
   "source": [
    "utils.val_scores_before_after_train(\n",
    "    cora.y[cora.val_mask],\n",
    "    yhat_naive_before_training[cora.val_mask],\n",
    "    yhat_naive_after_training[cora.val_mask],\n",
    "    \"Naiwny GAT\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 43328 characters]",
      "text/plain": [
       "<Figure size 472.441x472.441 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "utils.plot_results(\n",
    "    epoch_losses_naive,\n",
    "    val_losses_naive,\n",
    "    plot_name=\"gat_naive.tiff\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----------------------------------+--------------------------+----------------+----------+\n",
      "| Layer                             | Input Shape              | Output Shape   | #Param   |\n",
      "|-----------------------------------+--------------------------+----------------+----------|\n",
      "| NodesClassifierModel              | [2708, 1433], [2, 10556] | [2708, 7]      | 209,675  |\n",
      "| ├─(gnn_convs)ModuleList           | --                       | --             | 200,964  |\n",
      "| │    └─(0)NaiveGAT                | [2708, 1433], [2, 10556] | [2708, 256]    | 184,194  |\n",
      "| │    │    └─(aggr)MeanAggregation | [4, 2, 128]              | [1, 2, 128]    | --       |\n",
      "| │    │    └─(sigma)Identity       | [1, 256]                 | [1, 256]       | --       |\n",
      "| │    │    └─(W)Linear             | [2708, 1433]             | [2708, 128]    | 183,424  |\n",
      "| │    │    └─(att)Linear           | [4, 256]                 | [4, 2]         | 514      |\n",
      "| │    └─(1)NaiveGAT                | [2708, 256], [2, 10556]  | [2708, 128]    | 16,770   |\n",
      "| │    │    └─(aggr)MeanAggregation | [4, 2, 128]              | [1, 2, 128]    | --       |\n",
      "| │    │    └─(sigma)Identity       | [1, 256]                 | [1, 256]       | --       |\n",
      "| │    │    └─(W)Linear             | [2708, 256]              | [2708, 64]     | 16,384   |\n",
      "| │    │    └─(att)Linear           | [4, 128]                 | [4, 2]         | 258      |\n",
      "| ├─(lin_layers)Sequential          | [2708, 128]              | [2708, 7]      | 8,711    |\n",
      "| │    └─(0)Linear                  | [2708, 128]              | [2708, 64]     | 8,256    |\n",
      "| │    └─(1)ReLU                    | [2708, 64]               | [2708, 64]     | --       |\n",
      "| │    └─(2)Linear                  | [2708, 64]               | [2708, 7]      | 455      |\n",
      "| │    └─(3)ReLU                    | [2708, 7]                | [2708, 7]      | --       |\n",
      "+-----------------------------------+--------------------------+----------------+----------+\n"
     ]
    }
   ],
   "source": [
    "print(pyg_nn.summary(naive_clf, cora.x, cora.edge_index))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model klasyfikacji z warstwą splotu GAT z PyG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyg.seed_everything(123)\n",
    "\n",
    "nheads = 2\n",
    "out1 = 128\n",
    "out2 = 64\n",
    "\n",
    "pyg_gat1 = pyg_nn.GATConv(in_channels=cora.num_features, out_channels=out1, aggr='mean', heads=nheads)\n",
    "pyg_gat2 = pyg_nn.GATConv(in_channels=out1*nheads, out_channels=out2, aggr='mean', heads=nheads)\n",
    "\n",
    "pyg_clf = NodesClassifierModel(gnn_convs=(pyg_gat1, pyg_gat2), lin_sizes=(out2*nheads, 64, 7))\n",
    "optim = th.optim.Adam(pyg_clf.parameters(), lr=0.01)\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "with th.no_grad():\n",
    "    pyg_clf.eval()\n",
    "    yhat_pyg_before_training = pyg_clf(cora.x, cora.edge_index)\n",
    "\n",
    "\n",
    "epoch_losses_pyg = []\n",
    "val_losses_pyg = []\n",
    "\n",
    "for epoch in tqdm(range(30)):\n",
    "    pyg_clf.train()\n",
    "    optim.zero_grad()\n",
    "    out = pyg_clf(cora.x, cora.edge_index)\n",
    "    loss = loss_fn(out[cora.train_mask], cora.y[cora.train_mask])\n",
    "    loss.backward()\n",
    "    optim.step()\n",
    "    epoch_losses_pyg.append(loss.item())\n",
    "\n",
    "    pyg_clf.eval()\n",
    "    with th.no_grad():\n",
    "        val_loss = loss_fn(out[cora.val_mask], cora.y[cora.val_mask])\n",
    "        val_losses_pyg.append(val_loss.item())\n",
    "\n",
    "with th.no_grad():\n",
    "    pyg_clf.eval()\n",
    "    yhat_pyg_after_training = pyg_clf(cora.x, cora.edge_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wynik klasyfikacji modelu PyG GAT przed trenowaniem:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.00      0.00      0.00        61\n",
      "           1       0.07      1.00      0.13        36\n",
      "           2       0.00      0.00      0.00        78\n",
      "           3       0.00      0.00      0.00       158\n",
      "           4       0.00      0.00      0.00        81\n",
      "           5       0.00      0.00      0.00        57\n",
      "           6       0.00      0.00      0.00        29\n",
      "\n",
      "    accuracy                           0.07       500\n",
      "   macro avg       0.01      0.14      0.02       500\n",
      "weighted avg       0.01      0.07      0.01       500\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "Wynik klasyfikacji modelu PyG GAT po trenowaniu:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.47      0.75      0.58        61\n",
      "           1       0.58      0.69      0.63        36\n",
      "           2       0.78      0.91      0.84        78\n",
      "           3       0.59      0.80      0.68       158\n",
      "           4       0.00      0.00      0.00        81\n",
      "           5       0.78      0.74      0.76        57\n",
      "           6       0.00      0.00      0.00        29\n",
      "\n",
      "    accuracy                           0.62       500\n",
      "   macro avg       0.46      0.56      0.50       500\n",
      "weighted avg       0.50      0.62      0.55       500\n",
      "\n"
     ]
    }
   ],
   "source": [
    "utils.val_scores_before_after_train(\n",
    "    cora.y[cora.val_mask],\n",
    "    yhat_pyg_before_training[cora.val_mask],\n",
    "    yhat_pyg_after_training[cora.val_mask],\n",
    "    \"PyG GAT\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 42356 characters]",
      "text/plain": [
       "<Figure size 472.441x472.441 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "utils.plot_results(\n",
    "    epoch_losses_pyg,\n",
    "    val_losses_pyg,\n",
    "    plot_name=\"rys_4.11.tiff\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------------+--------------------------+----------------+----------+\n",
      "| Layer                    | Input Shape              | Output Shape   | #Param   |\n",
      "|--------------------------+--------------------------+----------------+----------|\n",
      "| NodesClassifierModel     | [2708, 1433], [2, 10556] | [2708, 7]      | 409,479  |\n",
      "| ├─(gnn_convs)ModuleList  | --                       | --             | 400,768  |\n",
      "| │    └─(0)GATConv        | [2708, 1433], [2, 10556] | [2708, 256]    | 367,616  |\n",
      "| │    └─(1)GATConv        | [2708, 256], [2, 10556]  | [2708, 128]    | 33,152   |\n",
      "| ├─(lin_layers)Sequential | [2708, 128]              | [2708, 7]      | 8,711    |\n",
      "| │    └─(0)Linear         | [2708, 128]              | [2708, 64]     | 8,256    |\n",
      "| │    └─(1)ReLU           | [2708, 64]               | [2708, 64]     | --       |\n",
      "| │    └─(2)Linear         | [2708, 64]               | [2708, 7]      | 455      |\n",
      "| │    └─(3)ReLU           | [2708, 7]                | [2708, 7]      | --       |\n",
      "+--------------------------+--------------------------+----------------+----------+\n"
     ]
    }
   ],
   "source": [
    "print(pyg_nn.summary(pyg_clf, cora.x, cora.edge_index))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Alternatywna forma implementacji PyG\n",
    "\n",
    "Biblioteka PyG umożliwia tworzenie warstw splotu na dwa sposoby:\n",
    "1. Albo jako pojedyncze warstwy - każda z nich jest osobnym obiektem, który można wykorzystać w modelu (odmiennie ustawiając parametry itd.). Wtedy warstwy są łączone w modelu, tak jak w przypadku modelu z naiwną implementacją.\n",
    "2. Albo jako jeden model, który jest automatycznie konstruowany przez PyG z kilku warstw splotu zestawionych razem. \n",
    "\n",
    "\n",
    "W przedstawionym przykładzie zastosowaliśmy podejście (1) - samodzielnie opisując każdą warstwę. Poniżej przedstawiony jest przykład podejścia (2): PyG zbuduje jeden model grafowy, z dwiema warstwami GAT. Wejściem do każdej będzie:\n",
    "* macierz cech wierzchołków (wyjście z poprzedniej warstwy lub oryginalne cechy wierzchołków),\n",
    "* macierz sąsiedztwa grafu (nie ulega zmianie).\n",
    "\n",
    "Określamy specyficzne dla GAT parametry, które będą miały zastosowanie **do wszystkich warstw**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------------------+--------------------------+----------------+----------+\n",
      "| Layer               | Input Shape              | Output Shape   | #Param   |\n",
      "|---------------------+--------------------------+----------------+----------|\n",
      "| GAT                 | [2708, 1433], [2, 10556] | [2708, 64]     | 400,576  |\n",
      "| ├─(dropout)Dropout  | [2708, 128]              | [2708, 128]    | --       |\n",
      "| ├─(act)ReLU         | [2708, 128]              | [2708, 128]    | --       |\n",
      "| ├─(convs)ModuleList | --                       | --             | 400,576  |\n",
      "| │    └─(0)GATv2Conv | [2708, 1433], [2, 10556] | [2708, 128]    | 367,360  |\n",
      "| │    └─(1)GATv2Conv | [2708, 128], [2, 10556]  | [2708, 64]     | 33,216   |\n",
      "| ├─(norms)ModuleList | --                       | --             | --       |\n",
      "| │    └─(0)Identity  | [2708, 128]              | [2708, 128]    | --       |\n",
      "| │    └─(1)Identity  | --                       | --             | --       |\n",
      "+---------------------+--------------------------+----------------+----------+\n"
     ]
    }
   ],
   "source": [
    "gat_model = pyg_nn.GAT(\n",
    "    in_channels=cora.num_features, # Ilość cech wejściowych - oryginalnych\n",
    "    hidden_channels=128,           # Ilość cech w warstwie ukrytej\n",
    "    num_layers=2,                  # Ilość warstw ukrytych - każda warstwa jest typu GATConv\n",
    "    heads=2,                       # Ilość głów uwagi\n",
    "    out_channels=64,               # Ilość klas do przewidzenia\n",
    "    # Dodatkowe hiperparametry dla każdej warstwy\n",
    "    aggr='mean',\n",
    "    v2=True\n",
    ")\n",
    "\n",
    "\n",
    "print(pyg_nn.summary(gat_model, cora.x, cora.edge_index))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Taki model można używać identycznie, jak wcześniej."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "with th.no_grad():\n",
    "    yhat = gat_model(cora.x, cora.edge_index)\n",
    "    classifier = NodesClassifierModel(gnn_convs=(gat_model,), lin_sizes=(64, 32, 7))\n",
    "    yhat_classifier = classifier(cora.x, cora.edge_index)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
