labs/lab10-gnn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9f0736f9-7034-4c07-b88f-62a8c18dfe93",
   "metadata": {},
   "source": [
    "# Grafové neurónové siete\n",
    "\n",
    "Grafové neurónové siete sú modely, ktoré pracujú s grafovými dátami a môžu byť využité okrem iného na klasifikáciu textu, v odporúčacích systémoch, v predpovedaní premávky a pri vývoji nových liekov.\n",
    "\n",
    "Grafy sú abstraktné a nelineárne údajové štruktúry, ktoré vyjadrujú vzťah medzi rôznymi jednotkami údajov. Môžu byť orientované alebo neorientované, a vážené alebo nevážené (v uzloch, v hranách a v oboch). Graf je vo všetkých týchto prípadoch reprezentovaný cez maticu susednosti, kde každá hodnota je 0 alebo 1, vyjadrujúca existenciu resp. neexistenciu spojenia medzi dvomi uzlami. Pri neorientovaných grafoch platí, že $x_{ij} = x_{ji}$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b82fd2da-8fe5-4572-bbba-24e6e0baf068",
   "metadata": {},
   "source": [
    "![](lab10/molecule.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33de1f03-9b76-4f20-a742-ddd28d9a1edc",
   "metadata": {},
   "source": [
    "## Konvolúcia nad grafmi\n",
    "\n",
    "Podobne ako obrázky, aj grafy spĺňajú niekoľko kritérií, ktoré ich predurčujú na úspešné spracovanie konvolúciou. Aj pre nich platí lokalita, teda súvisiace údaje a informácie sú uložené zvyčajne blízko seba (vrchol vzdialený o jednu hranu bude relevantnejší ako vrchol vzdialený niekoľko hrán). Pre grafy navyše platí permutačná invariancia, teda môžeme pomiešať poradie uzlov, a reprezentácia grafu sa nemení. Konvolúcia dokáže úspešne pracovať s takouto invarianciou, a počet parametrov stále nezávisí od veľkosti vstupu. V skutočnosti existuje niekoľko konvolučných prístupov pre spracovanie grafových dát, a ďalšie sú stále predmetom výskumu."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8fa2d3f-4fca-4931-a92a-585db2e0d418",
   "metadata": {},
   "source": [
    "Majme nasledovný graf:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bd50c62-25b5-442d-b80f-3fbad550ff51",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import networkx as nx\n",
    "\n",
    "G = nx.Graph()\n",
    "blue, orange, green = \"#1f77b4\", \"#ff7f0e\", \"#2ca02c\"\n",
    "G.add_nodes_from([\n",
    "    (1, {\"color\": blue}),\n",
    "    (2, {\"color\": orange}),\n",
    "    (3, {\"color\": blue}),\n",
    "    (4, {\"color\": green})\n",
    "])\n",
    "G.add_edges_from([(1,2), (2,3), (1,3), (3,4)])\n",
    "A = np.asarray(nx.adjacency_matrix(G).todense())\n",
    "print(A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d31018a-7b35-42d5-979d-d8f3b16f2b0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_graph_color_label_representation(G, mapping_dict):\n",
    "    one_hot_idxs = np.array([mapping_dict[v] for v in\n",
    "    nx.get_node_attributes(G, 'color').values()])\n",
    "\n",
    "    one_hot_encoding = np.zeros((one_hot_idxs.size, len(mapping_dict)))\n",
    "    one_hot_encoding[np.arange(one_hot_idxs.size), one_hot_idxs] = 1\n",
    "    return one_hot_encoding\n",
    "\n",
    "X = build_graph_color_label_representation(G, {green: 0, blue: 1, orange: 2})\n",
    "print(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c37256e-5df1-4082-98c0-2e021870f546",
   "metadata": {},
   "outputs": [],
   "source": [
    "color_map = nx.get_node_attributes(G, 'color').values()\n",
    "nx.draw(G,with_labels=True, node_color=color_map)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "846dfec8-af5b-4618-a451-c50e8edc85d1",
   "metadata": {},
   "source": [
    "Grafová konvolúcia aktualizuje embedding jednotlivých uzlov (každému prislúcha jeden riadok v matici) s informáciami o samotnom uzle a jeho susedoch. Konvolúcia môže prebiehať napríklad:\n",
    "\n",
    "$$x'_{i} = x_i \\cdot W_1 + \\sum_{j \\in N(i)} x_j \\cdot W_2 + b,$$\n",
    "\n",
    "kde $W_1$ a $W_2$ sú trénovateľné váhy (s rozmermi $f_{in} \\times f_{out}$) a $b$ je trénovateľný bias (s rozmermi $f_{out}$). Úlohou váh je slúžiť ako filtre a berú do úvahy iba priame spojenia. Pre zohľadnenie vzťahov o niekoľko hrán musíme použiť viacnásobnú konvolúciu. Potom pomocou matice susednosti vieme vypočítať požadované hodnoty:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "650a5ea0-c40e-47a9-956b-925cc1c394cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "f_in, f_out = X.shape[1], 6\n",
    "W_1 = np.random.rand(f_in, f_out)\n",
    "W_2 = np.random.rand(f_in, f_out)\n",
    "h = np.dot(X, W_1) + np.dot(np.dot(A, X), W_2)\n",
    "print(h)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09c7b820-a915-46e2-8c36-300aa356d17a",
   "metadata": {},
   "source": [
    "Rôzne grafové konvolúcie majú jednotnú schému spracovania cez posielanie správ, kde každý uzol má svoj vlastný skrytý stav $h_i^{(t)}$, pričom $h_i^{(0)} = x_i$ (riadok pre uzol *i*). Následne sa konvolúcia vykoná v dvoch fázach:\n",
    "\n",
    "1. fáza posielania správ - pre každý uzol sa vypočíta hodnota funkcie správy (*message function*):\n",
    "$$m_i = \\sum_{j \\in N(i)} M_t(h_i^{(t)}, h_j^{(t)}, e_{ij}),$$\n",
    "kde $M_t$ je funkcia správy.\n",
    "2. fáza aktualizácie vrcholov - skrytý stav vrcholov sa aktualizuje podľa:\n",
    "$$h_i^{(t+1)} = U_t(h_i^{(t)}, m_i)$$\n",
    "\n",
    "**Poznámka:** Pre orientované grafy sa susednosť definuje cez hrany smerujúce do uzla *i*."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc2af715-51bf-4620-8483-f531e71c85bd",
   "metadata": {},
   "source": [
    "## Definícia grafovej siete v PyTorchi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2044d6e-49e1-4b0b-83cf-e3d393b55cbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import networkx as nx\n",
    "import torch\n",
    "from torch.nn.parameter import Parameter\n",
    "import numpy as np\n",
    "import math\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class NodeNetwork(torch.nn.Module):\n",
    "    def __init__(self, input_features):\n",
    "        super().__init__()\n",
    "        # network with two graph convolutions\n",
    "        self.conv_1 = BasicGraphConvolutionLayer(input_features, 32)\n",
    "        self.conv_2 = BasicGraphConvolutionLayer(32, 32)\n",
    "        self.fc_1 = torch.nn.Linear(32, 16)\n",
    "        self.out_layer = torch.nn.Linear(16, 2)\n",
    "\n",
    "    def forward(self, X, A, batch_mat):\n",
    "        x = F.relu(self.conv_1(X, A))\n",
    "        x = F.relu(self.conv_2(x, A))\n",
    "        # use global pooling for dealing with graphs of varying sizes\n",
    "        output = global_sum_pool(x, batch_mat)\n",
    "        output = self.fc_1(output)\n",
    "        output = self.out_layer(output)\n",
    "        return F.softmax(output, dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63ef49dd-ca57-472a-954c-c07e0f66990e",
   "metadata": {},
   "source": [
    "Štruktúra siete je vizualizovaná na nasledujúcom obrázku:\n",
    "\n",
    "![](lab10/structure.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b9dfb68-7053-4050-982c-02c2141a7be2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BasicGraphConvolutionLayer(torch.nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super().__init__()\n",
    "        self.in_channels = in_channels\n",
    "        self.out_channels = out_channels\n",
    "        # generate random weights and biases\n",
    "        self.W2 = Parameter(torch.rand(\n",
    "            (in_channels, out_channels), dtype=torch.float32))\n",
    "        self.W1 = Parameter(torch.rand(\n",
    "            (in_channels, out_channels), dtype=torch.float32))\n",
    "        self.bias = Parameter(torch.zeros(\n",
    "            out_channels, dtype=torch.float32))\n",
    "\n",
    "    def forward(self, X, A):\n",
    "        potential_msgs = torch.mm(X, self.W2)\n",
    "        propagated_msgs = torch.mm(A, potential_msgs)\n",
    "        root_update = torch.mm(X, self.W1)\n",
    "        output = root_update + propagated_msgs + self.bias\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "363e32d7-2e9f-4131-8e01-cc1bca9c728a",
   "metadata": {},
   "source": [
    "Príklad použitia grafovej konvolúcie:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d96a1d7-6920-42bd-a755-67c413717a1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('X.shape:', X.shape)\n",
    "print('A.shape:', A.shape)\n",
    "\n",
    "basiclayer = BasicGraphConvolutionLayer(3, 8)\n",
    "out = basiclayer(\n",
    "    X=torch.tensor(X, dtype=torch.float32),\n",
    "    A=torch.tensor(A, dtype=torch.float32)\n",
    ")\n",
    "print('Output shape:', out.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0f164d9-27be-49d0-98b1-8c35983698b8",
   "metadata": {},
   "source": [
    "Cieľom globálneho *pooling*u je agregovať všetky uzlové embeddingy grafu do jednotnej veľkosti (podobne ako pri spracovaní obrazu cez global average pooling). Pooling môže fungovať na základe funkcií, ktoré sú invariantné voči permutáciám (napríklad suma, maximum, priemer):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1493dfa-0b36-4c82-aace-5f1235b228cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def global_sum_pool(X, batch_mat):\n",
    "    if batch_mat is None or batch_mat.dim() == 1:\n",
    "        return torch.sum(X, dim=0).unsqueeze(0)\n",
    "    else:\n",
    "        return torch.mm(batch_mat, X)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4b311ef-665e-4888-83a8-b1fde7252534",
   "metadata": {},
   "source": [
    "Vytváranie mini-batchov pri grafoch nie je také priamočiare ako v prípade obrázkov, najmä vzhľadom na to, že grafy môžu mať rôzne veľkosti. Tento problém sa dá obísť paddingom, čo ale stále nefunguje ideálne v prípade, že veľkosť grafov sa líši vo väčšej miere. Práve preto mini-batch je často reprezentovaný ako jeden graf s niekoľkými nespojenými podgrafmi."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0b6e761-8288-45cc-8bf8-cf0f1a796f06",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_batch_tensor(graph_sizes):\n",
    "    starts = [sum(graph_sizes[:idx])\n",
    "              for idx in range(len(graph_sizes))]\n",
    "    stops = [starts[idx] + graph_sizes[idx]\n",
    "             for idx in range(len(graph_sizes))]\n",
    "    tot_len = sum(graph_sizes)\n",
    "    batch_size = len(graph_sizes)\n",
    "    batch_mat = torch.zeros([batch_size, tot_len]).float()\n",
    "    for idx, starts_and_stops in enumerate(zip(starts, stops)):\n",
    "        start = starts_and_stops[0]\n",
    "        stop = starts_and_stops[1]\n",
    "        batch_mat[idx,start:stop] = 1\n",
    "    return batch_mat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca011455-0a3e-4147-a1d3-837614340dc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# batch is a list of dictionaries each containing\n",
    "# the representation and label of a graph\n",
    "def collate_graphs(batch):\n",
    "    adj_mats = [graph['A'] for graph in batch]\n",
    "    sizes = [A.size(0) for A in adj_mats]\n",
    "    tot_size = sum(sizes)\n",
    "    # create batch matrix\n",
    "    batch_mat = get_batch_tensor(sizes)\n",
    "    # combine feature matrices\n",
    "    feat_mats = torch.cat([graph['X'] for graph in batch], dim=0)\n",
    "    # combine labels\n",
    "    labels = torch.cat([graph['y'] for graph in batch], dim=0)\n",
    "    # combine adjacency matrices\n",
    "    batch_adj = torch.zeros([tot_size, tot_size], dtype=torch.float32)\n",
    "    accum = 0\n",
    "    for adj in adj_mats:\n",
    "        g_size = adj.shape[0]\n",
    "        batch_adj[accum:accum+g_size,accum:accum+g_size] = adj\n",
    "        accum = accum + g_size\n",
    "    repr_and_label = {'A': batch_adj,\n",
    "                      'X': feat_mats, 'y': labels,\n",
    "                      'batch': batch_mat}\n",
    "    return repr_and_label"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec46e3ef-26bf-4512-91c8-c0c715c4e1b6",
   "metadata": {},
   "source": [
    "## Vytvorenie grafov a príprava údajov"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84869de5-4ccd-405d-832d-4644b7b8ddd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_graph_dict(G, mapping_dict):\n",
    "    # Function builds dictionary representation of graph G\n",
    "    A = torch.from_numpy(\n",
    "        np.asarray(nx.adjacency_matrix(G).todense())).float()\n",
    "    # build_graph_color_label_representation()\n",
    "    # was introduced with the first example graph\n",
    "    X = torch.from_numpy(\n",
    "        build_graph_color_label_representation(\n",
    "            G, mapping_dict)).float()\n",
    "    # kludge since there is not specific task for this example\n",
    "    y = torch.tensor([[1,0]]).float()\n",
    "    return {'A': A, 'X': X, 'y': y, 'batch': None}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cd9e5e3-320c-4ed2-acad-42767400e2f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "blue, orange, green = \"#1f77b4\", \"#ff7f0e\",\"#2ca02c\"\n",
    "mapping_dict= {green:0, blue:1, orange:2}\n",
    "G1 = nx.Graph()\n",
    "G1.add_nodes_from([\n",
    "    (1,{\"color\": blue}),\n",
    "    (2,{\"color\": orange}),\n",
    "    (3,{\"color\": blue}),\n",
    "    (4,{\"color\": green})\n",
    "])\n",
    "G1.add_edges_from([(1, 2), (2, 3), (1, 3), (3, 4)])\n",
    "\n",
    "G2 = nx.Graph()\n",
    "G2.add_nodes_from([\n",
    "    (1,{\"color\": green}),\n",
    "    (2,{\"color\": green}),\n",
    "    (3,{\"color\": orange}),\n",
    "    (4,{\"color\": orange}),\n",
    "    (5,{\"color\": blue})\n",
    "])\n",
    "G2.add_edges_from([(2, 3),(3, 4),(3, 1),(5, 1)])\n",
    "\n",
    "G3 = nx.Graph()\n",
    "G3.add_nodes_from([\n",
    "    (1,{\"color\": orange}),\n",
    "    (2,{\"color\": orange}),\n",
    "    (3,{\"color\": green}),\n",
    "    (4,{\"color\": green}),\n",
    "    (5,{\"color\": blue}),\n",
    "    (6,{\"color\":orange})\n",
    "])\n",
    "G3.add_edges_from([(2,3), (3,4), (3,1), (5,1), (2,5), (6,1)])\n",
    "\n",
    "G4 = nx.Graph()\n",
    "G4.add_nodes_from([\n",
    "    (1,{\"color\": blue}),\n",
    "    (2,{\"color\": blue}),\n",
    "    (3,{\"color\": green})\n",
    "])\n",
    "G4.add_edges_from([(1, 2), (2, 3)])\n",
    "graph_list = [\n",
    "    get_graph_dict(graph, mapping_dict) for graph in [G1, G2, G3, G4]\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4855d6df-8017-4be2-bf1e-b8b04fcf1065",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset\n",
    "\n",
    "class ExampleDataset(Dataset):\n",
    "    # Simple PyTorch dataset that will use our list of graphs\n",
    "    def __init__(self, graph_list):\n",
    "        self.graphs = graph_list\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.graphs)\n",
    "\n",
    "    def __getitem__(self,idx):\n",
    "        mol_rep = self.graphs[idx]\n",
    "        return mol_rep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4f320a0-d3e6-4b00-b839-862ca9f8f387",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "dset = ExampleDataset(graph_list)\n",
    "loader = DataLoader(\n",
    "    dset, batch_size=2, shuffle=False,\n",
    "    collate_fn=collate_graphs\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eed93435-a538-4a79-a631-ee0bce7bd1e0",
   "metadata": {},
   "source": [
    "## Generovanie predikcií"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab499fdc-4dcd-4c71-a993-d68cb6c41d9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_features = 3\n",
    "net = NodeNetwork(node_features)\n",
    "batch_results = []\n",
    "for b in loader:\n",
    "    batch_results.append(net(b['X'], b['A'], b['batch']).detach())\n",
    "print(batch_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61ecc36d-926c-4920-b93a-d3fedcdb2471",
   "metadata": {},
   "outputs": [],
   "source": [
    "G2_rep = dset[1]\n",
    "G2_single = net(G2_rep['X'], G2_rep['A'], G2_rep['batch']).detach()\n",
    "print(G2_single)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1425e966-e9f2-423a-9248-bb4974a35eab",
   "metadata": {},
   "outputs": [],
   "source": [
    "G2_batch = batch_results[0][1]\n",
    "print(G2_batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d12fec04-bcba-4d8c-9842-b404fb881aa5",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.all(torch.isclose(G2_single, G2_batch))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8864db7-1b46-4fc9-9c54-930873f99389",
   "metadata": {},
   "source": [
    "## Príklady\n",
    "\n",
    "* [Definícia a trénovanie GNN pomocou PyTorch Geometric](https://github.com/rasbt/machine-learning-book/blob/main/ch18/ch18_part2.ipynb)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7751de62-7045-4a7b-842a-c10d4e3634c2",
   "metadata": {},
   "source": [
    "## Použité zdroje\n",
    "\n",
    "* Raschka, Sebastian, Yuxi Hayden Liu, Vahid Mirjalili, and Dmytro Dzhulgakov. Machine Learning with PyTorch and Scikit-Learn: Develop machine learning and deep learning models with Python. Packt Publishing Ltd, 2022. Kapitola 18\n",
    "* Dehmamy, Nima, Albert-László Barabási, and Rose Yu. \"Understanding the representation power of graph neural networks in learning graph topology.\" Advances in Neural Information Processing Systems 32 (2019).\n",
    "* Gilmer, Justin, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl. \"Neural message passing for quantum chemistry.\" In International conference on machine learning, pp. 1263-1272. PMLR, 2017."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}