{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wprowadzenie\n",
    "\n",
    "W tym notatniku przyjrzymy się splotowi SAGE. Publikacja, w której został zaprezentowany wywarła bardzo silny wpływ na kolejne implementacje, głównie ze względu na poruszenie bardzo istotnego problemu próbkowania w grafach oraz metodyki prowadzenia szkolenia sieci.\n",
    "\n",
    "Oryginalna publikacja:\n",
    "\n",
    "> Hamilton, W., Ying, Z., & Leskovec, J. (2017). Inductive representation learning on large graphs. Advances in neural information processing systems, 30."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import bibliotek"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\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",
    "\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": [
    "# SAGE\n",
    "\n",
    "Warstwa splotu SAGE została przedstawiona wspólnie z zarysem koncepcji prowadzenia treningu sieci grafowych. Ten ostatni aspekt zostanie szczegółowo omówiony w kolejnych rozdziałach. Poniżej skupimy się na samej implementacji warstwy splotu. Autorzy proponują trzy możliwe warianty agregacji dla tej warstwy:\n",
    "1. Uśrednianie wektorów osadzenia sąsiedztwa - takie podejście w założeniach teoretycznych ma upodabniać SAGE do warstwy GCN. Dodatkowo, formalizacja dla uśredniania zawiera w sobie formę *połączenia skrótowego* (ang. \"skip-connection\") - dodając reprezentację własną wierzchołka v.\n",
    "2. Uśrednianie max - inna forma agregacji, opierająca się na początkowym przekształceniu wektorów osadzenia sąsiedztwa macierzą wag, a potem wybraniu maksymalnej wartości z każdej kolumny.\n",
    "3. Agregacja LSTM - dodatkowa forma agregacji, która nie posiada własności symetrii (w przeciwieństwie do np. średniej), ale pozwala na bardziej złożone przekształcenia.\n",
    "\n",
    "Poniżej przedstawiona jest najbardziej ogólna forma warstwy SAGE z oryginalnej publikacji:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\begin{align*}\n",
    "\\mathbf{h}_{\\mathcal{N}(v)}^{(k)} &= \\text{AGGREGATE}\\left\\{\\mathbf{h}_u^{(k-1)}, \\forall u \\in \\mathcal{N}(v) \\right\\} & \\triangleright \\text{Agregacja sąsiadów} \\\\\n",
    "\\mathbf{h}_v^{(k)} &= \\sigma\\left(\\mathbf{W}^{(k)} \\cdot \\text{CONCAT}(\\mathbf{h}_v^{(k-1)}, \\mathbf{h}_{\\mathcal{N}(v)}^{(k)})\\right) & \\triangleright \\text{Aktualizacja reprezentacji wierzchołka}\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "gdzie:\n",
    "\n",
    "- $\\mathbf{h}_v^{(k)}$ to reprezentacja wierzchołka $v$ w warstwie $k$;\n",
    "- $\\mathbf{W}^{(k)}$ to macierz wag warstwy;\n",
    "- $\\mathcal{N}(v)$ - zbiór sąsiadów wierzchołka $v$;\n",
    "- $\\sigma$ funkcja aktywacji - zwykle nieliniowa (jak np. ReLU). Można ją opuścić, jeśli sieć GNN jest częścią większej architektury."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatywna wersja tej formalizacji, stosowana np. w bibliotece PyG, zakłada nieco zmienioną formę obu równań. Jest ona bardzo zbliżona do tej, sugerowanej przez autorów publikacji, jako domyślną, dla agregacji uśredniającej.\n",
    "\n",
    "1. Osobny zestaw wag do przekształcenia zagregowanych sąsiadów.\n",
    "2. Drugi zestaw wag - przeznaczony tylko dla wierzchołka *v*."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\begin{align*}\n",
    "\\mathbf{h}_{\\mathcal{N}(v)}^{(k)} &= \\mathbf{W}_2 \\cdot \\text{ŚREDNIA}\\left\\{\\mathbf{h}_u^{(k-1)}, \\forall u \\in \\mathcal{N}(v) \\right\\} & \\triangleright \\text{Agregacja sąsiadów} \\\\\n",
    "\\mathbf{h}_v^{(k)} &= \\mathbf{W}_1 \\mathbf{h}_v^{(k-1)} + \\mathbf{h}_{\\mathcal{N}(v)}^{(k)} & \\triangleright \\text{Aktualizacja reprezentacji wierzchołka}\n",
    "\\end{align*}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Klasyfikacja wierzchołków\n",
    "\n",
    "W naiwnej implementacji uwzględnimy podejście przeznaczone na agregacji uśredniającej."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementacja naiwna"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NaiveSAGE(NaiveGraphConv):\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            dim_in: int,\n",
    "            dim_out: int,\n",
    "            project: bool = True,\n",
    "            normalize: bool = True,\n",
    "            sigma=nn.Identity(),\n",
    "            agg_fn=pyg_nn.MeanAggregation()):\n",
    "        super().__init__(dim_in, dim_out, aggr=agg_fn, sigma=sigma, add_self_loops=False)\n",
    "        self.project = project\n",
    "        self.W_project = nn.Linear(dim_in, dim_in, bias=True)\n",
    "        self.W_agg = nn.Linear(dim_in, dim_out, bias=True)\n",
    "        self.W_self = nn.Linear(dim_in, dim_out, bias=False)\n",
    "        self.normalize = normalize\n",
    "        self.reset_parameters()\n",
    "\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        pyg_nn.inits.glorot(self.W_agg.weight)\n",
    "        pyg_nn.inits.glorot(self.W_project.weight)\n",
    "        pyg_nn.inits.glorot(self.W_self.weight)\n",
    "        self.aggr.reset_parameters()\n",
    "\n",
    "    def message(self, v_neighbors: 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 - nie zachodzą przekształcenia wiadomości.\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",
    "        Returns\n",
    "        -------\n",
    "        th.Tensor\n",
    "            Macierz wiadomości od sąsiadów.\n",
    "            Wymiar: `|N(v)| x dim_out`.\n",
    "        \"\"\"\n",
    "        if self.project:\n",
    "            return self.W_project(v_neighbors).relu()\n",
    "        else:\n",
    "            return v_neighbors\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 dim_out`.\n",
    "        \"\"\"\n",
    "        # Po agregacji upewniamy się, że wynikowy wektor zachowa wymiar (1 x dim out).\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 SAGE - wykorzystuje uśrednione wartości\n",
    "        wiadomości sąsiadów oraz własną reprezentację wierzchołka v.\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",
    "        v_new = th.atleast_2d(self.W_self(v))\n",
    "        agg_new = self.W_agg(agg)\n",
    "        new_repr = self.sigma(th.cat([v_new, agg_new], dim=-1))\n",
    "        return new_repr\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.\n",
    "\n",
    "        edge_index : th.Tensor\n",
    "            Macierz sąsiedztwa grafu.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        th.Tensor\n",
    "            Zaktualizowane cechy wierzchołków.\n",
    "        \"\"\"\n",
    "        # Alokacja nowej macierzy reprezentacji wierzchołków\n",
    "        h_new = super().forward(x, edge_index)\n",
    "        if self.normalize:\n",
    "            return nn.functional.normalize(h_new, p=2., dim=-1)\n",
    "        else:\n",
    "            return h_new"
   ]
  },
  {
   "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_sage = NaiveSAGE(dim_in=cora.num_features, dim_out=16, project=True)\n",
    "    naive_sage.eval()\n",
    "    h = naive_sage(cora.x, cora.edge_index)\n",
    "h.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model klasyfikacji z naiwną implementacją SAGE"
   ]
  },
  {
   "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",
    "\n",
    "naive_sage1 = NaiveSAGE(\n",
    "    dim_in=cora.num_features, \n",
    "    dim_out=64, \n",
    "    project=True,\n",
    "    normalize=True,\n",
    "    sigma=nn.Linear(128, 64))\n",
    "naive_sage2 = NaiveSAGE(\n",
    "    dim_in=64, \n",
    "    dim_out=32,\n",
    "    project=True, \n",
    "    normalize=True, \n",
    "    sigma=nn.Linear(64, 32))\n",
    "\n",
    "naive_clf = NodesClassifierModel(gnn_convs=(naive_sage1, naive_sage2), lin_sizes=(32, 16, 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",
    "    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": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----------------------------------+--------------------------+----------------+-----------+\n",
      "| Layer                             | Input Shape              | Output Shape   | #Param    |\n",
      "|-----------------------------------+--------------------------+----------------+-----------|\n",
      "| NodesClassifierModel              | [2708, 1433], [2, 10556] | [2708, 7]      | 2,257,681 |\n",
      "| ├─(gnn_convs)ModuleList           | --                       | --             | 2,257,034 |\n",
      "| │    └─(0)NaiveSAGE               | [2708, 1433], [2, 10556] | [2708, 64]     | 2,246,666 |\n",
      "| │    │    └─(aggr)MeanAggregation | [3, 1433]                | [1, 1433]      | --        |\n",
      "| │    │    └─(sigma)Linear         | [1, 128]                 | [1, 64]        | 8,256     |\n",
      "| │    │    └─(W_project)Linear     | [3, 1433]                | [3, 1433]      | 2,054,922 |\n",
      "| │    │    └─(W_agg)Linear         | [1, 1433]                | [1, 64]        | 91,776    |\n",
      "| │    │    └─(W_self)Linear        | [1433]                   | [64]           | 91,712    |\n",
      "| │    └─(1)NaiveSAGE               | [2708, 64], [2, 10556]   | [2708, 32]     | 10,368    |\n",
      "| │    │    └─(aggr)MeanAggregation | [3, 1433]                | [1, 1433]      | --        |\n",
      "| │    │    └─(sigma)Linear         | [1, 64]                  | [1, 32]        | 2,080     |\n",
      "| │    │    └─(W_project)Linear     | [3, 64]                  | [3, 64]        | 4,160     |\n",
      "| │    │    └─(W_agg)Linear         | [1, 64]                  | [1, 32]        | 2,080     |\n",
      "| │    │    └─(W_self)Linear        | [64]                     | [32]           | 2,048     |\n",
      "| ├─(lin_layers)Sequential          | [2708, 32]               | [2708, 7]      | 647       |\n",
      "| │    └─(0)Linear                  | [2708, 32]               | [2708, 16]     | 528       |\n",
      "| │    └─(1)ReLU                    | [2708, 16]               | [2708, 16]     | --        |\n",
      "| │    └─(2)Linear                  | [2708, 16]               | [2708, 7]      | 119       |\n",
      "| │    └─(3)ReLU                    | [2708, 7]                | [2708, 7]      | --        |\n",
      "+-----------------------------------+--------------------------+----------------+-----------+\n"
     ]
    }
   ],
   "source": [
    "print(pyg_nn.summary(naive_clf, cora.x, cora.edge_index))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wynik klasyfikacji modelu Naiwny SAGE 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.32      1.00      0.48       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.32       500\n",
      "   macro avg       0.05      0.14      0.07       500\n",
      "weighted avg       0.10      0.32      0.15       500\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "Wynik klasyfikacji modelu Naiwny SAGE po trenowaniu:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.00      0.00      0.00        61\n",
      "           1       0.22      0.94      0.35        36\n",
      "           2       0.92      0.74      0.82        78\n",
      "           3       0.88      0.55      0.68       158\n",
      "           4       0.63      0.85      0.72        81\n",
      "           5       0.59      0.74      0.66        57\n",
      "           6       0.00      0.00      0.00        29\n",
      "\n",
      "    accuracy                           0.58       500\n",
      "   macro avg       0.46      0.55      0.46       500\n",
      "weighted avg       0.61      0.58      0.56       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 SAGE\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 42560 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=\"naive_sage.tiff\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model klasyfikacji z warstwą splotu SAGE z PyG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyg.seed_everything(123)\n",
    "\n",
    "\n",
    "pyg_sage1 = pyg_nn.SAGEConv(in_channels=cora.num_features, out_channels=64, project=True, normalize=True, aggr='mean')\n",
    "pyg_sage2 = pyg_nn.SAGEConv(in_channels=64, out_channels=32, project=True, normalize=True, aggr='mean')\n",
    "\n",
    "pyg_clf = NodesClassifierModel(gnn_convs=(pyg_sage1, pyg_sage2), lin_sizes=(32, 16, 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(25)):\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": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------------+--------------------------+----------------+-----------+\n",
      "| Layer                    | Input Shape              | Output Shape   | #Param    |\n",
      "|--------------------------+--------------------------+----------------+-----------|\n",
      "| NodesClassifierModel     | [2708, 1433], [2, 10556] | [2708, 7]      | 2,247,345 |\n",
      "| ├─(gnn_convs)ModuleList  | --                       | --             | 2,246,698 |\n",
      "| │    └─(0)SAGEConv       | [2708, 1433], [2, 10556] | [2708, 64]     | 2,238,410 |\n",
      "| │    └─(1)SAGEConv       | [2708, 64], [2, 10556]   | [2708, 32]     | 8,288     |\n",
      "| ├─(lin_layers)Sequential | [2708, 32]               | [2708, 7]      | 647       |\n",
      "| │    └─(0)Linear         | [2708, 32]               | [2708, 16]     | 528       |\n",
      "| │    └─(1)ReLU           | [2708, 16]               | [2708, 16]     | --        |\n",
      "| │    └─(2)Linear         | [2708, 16]               | [2708, 7]      | 119       |\n",
      "| │    └─(3)ReLU           | [2708, 7]                | [2708, 7]      | --        |\n",
      "+--------------------------+--------------------------+----------------+-----------+\n"
     ]
    }
   ],
   "source": [
    "print(pyg_nn.summary(pyg_clf, cora.x, cora.edge_index))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wynik klasyfikacji modelu PyG SAGE 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.16      1.00      0.27        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.16       500\n",
      "   macro avg       0.02      0.14      0.04       500\n",
      "weighted avg       0.02      0.16      0.04       500\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "Wynik klasyfikacji modelu PyG SAGE po trenowaniu:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.00      0.00      0.00        61\n",
      "           1       0.61      0.75      0.68        36\n",
      "           2       0.83      0.71      0.76        78\n",
      "           3       0.75      0.59      0.66       158\n",
      "           4       0.47      0.72      0.57        81\n",
      "           5       0.59      0.72      0.65        57\n",
      "           6       0.33      0.83      0.47        29\n",
      "\n",
      "    accuracy                           0.60       500\n",
      "   macro avg       0.51      0.62      0.54       500\n",
      "weighted avg       0.57      0.60      0.57       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 SAGE\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 38476 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.5.tiff\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Alternatywna forma implementacji PyG\n",
    "\n",
    "Biblioteka PyG umożliwia tworzenie warstw splotu na dwa sposoby:\n",
    "\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 SAGE. 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 GraphSAGE parametry, które będą miały zastosowanie **do wszystkich warstw**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------------------+--------------------------+----------------+-----------+\n",
      "| Layer               | Input Shape              | Output Shape   | #Param    |\n",
      "|---------------------+--------------------------+----------------+-----------|\n",
      "| GraphSAGE           | [2708, 1433], [2, 10556] | [2708, 32]     | 2,246,698 |\n",
      "| ├─(dropout)Dropout  | [2708, 64]               | [2708, 64]     | --        |\n",
      "| ├─(act)ReLU         | [2708, 64]               | [2708, 64]     | --        |\n",
      "| ├─(convs)ModuleList | --                       | --             | 2,246,698 |\n",
      "| │    └─(0)SAGEConv  | [2708, 1433], [2, 10556] | [2708, 64]     | 2,238,410 |\n",
      "| │    └─(1)SAGEConv  | [2708, 64], [2, 10556]   | [2708, 32]     | 8,288     |\n",
      "| ├─(norms)ModuleList | --                       | --             | --        |\n",
      "| │    └─(0)Identity  | [2708, 64]               | [2708, 64]     | --        |\n",
      "| │    └─(1)Identity  | --                       | --             | --        |\n",
      "+---------------------+--------------------------+----------------+-----------+\n"
     ]
    }
   ],
   "source": [
    "sage_model = pyg_nn.GraphSAGE(\n",
    "    in_channels=cora.num_features,  # Ilość cech wejściowych - oryginalnych\n",
    "    hidden_channels=64,             # Ilość cech w warstwie ukrytej\n",
    "    num_layers=2,                   # Ilość warstw ukrytych - każda warstwa jest warstwą typu SAGE\n",
    "    out_channels=32,                # Ilość klas do przewidzenia\n",
    "    # Dodatkowe hiperparametry dla każdej warstwy\n",
    "    normalize=True,\n",
    "    project=True,\n",
    "    aggr='mean'\n",
    ")\n",
    "\n",
    "\n",
    "print(pyg_nn.summary(sage_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": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "with th.no_grad():\n",
    "    yhat = sage_model(cora.x, cora.edge_index)\n",
    "    classifier = NodesClassifierModel(gnn_convs=(sage_model,), lin_sizes=(32, 16, 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
}
