docs/tutorial_notebooks/tutorial7/GNN_overview.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 7: Graph Neural Networks\n",
    "\n",
    "![Status](https://img.shields.io/static/v1.svg?label=Status&message=Finished&color=green)\n",
    "\n",
    "**Filled notebook:** \n",
    "[![View on Github](https://img.shields.io/static/v1.svg?logo=github&label=Repo&message=View%20On%20Github&color=lightgrey)](https://github.com/phlippe/uvadlc_notebooks/blob/master/docs/tutorial_notebooks/tutorial7/GNN_overview.ipynb)\n",
    "[![Open In Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/phlippe/uvadlc_notebooks/blob/master/docs/tutorial_notebooks/tutorial7/GNN_overview.ipynb)  \n",
    "**Pre-trained models:** \n",
    "[![View files on Github](https://img.shields.io/static/v1.svg?logo=github&label=Repo&message=View%20On%20Github&color=lightgrey)](https://github.com/phlippe/saved_models/tree/main/tutorial7)\n",
    "[![GoogleDrive](https://img.shields.io/static/v1.svg?logo=google-drive&logoColor=yellow&label=GDrive&message=Download&color=yellow)](https://drive.google.com/drive/folders/1DOTV_oYt5boa-MElbc2izat4VMSc1gob?usp=sharing)   \n",
    "**Recordings:** \n",
    "[![YouTube - Part 1](https://img.shields.io/static/v1.svg?logo=youtube&label=YouTube&message=Part%201&color=red)](https://youtu.be/fK7d56Ly9q8)\n",
    "[![YouTube - Part 2](https://img.shields.io/static/v1.svg?logo=youtube&label=YouTube&message=Part%202&color=red)](https://youtu.be/ZCNSUWe4a_Q)   \n",
    "**JAX+Flax version:**\n",
    "[![View on RTD](https://img.shields.io/static/v1.svg?logo=readthedocs&label=RTD&message=View%20On%20RTD&color=8CA1AF)](https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/JAX/tutorial7/GNN_overview.html)   \n",
    "**Author:** Phillip Lippe"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "**Note:** Interested in JAX? Check out our [JAX+Flax version](https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/JAX/tutorial7/GNN_overview.html) of this tutorial!\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial, we will discuss the application of neural networks on graphs. Graph Neural Networks (GNNs) have recently gained increasing popularity in both applications and research, including domains such as social networks, knowledge graphs, recommender systems, and bioinformatics. While the theory and math behind GNNs might first seem complicated, the implementation of those models is quite simple and helps in understanding the methodology. Therefore, we will discuss the implementation of basic network layers of a GNN, namely graph convolutions, and attention layers. Finally, we will apply a GNN on a node-level, edge-level, and graph-level tasks.\n",
    "\n",
    "Below, we will start by importing our standard libraries. We will use PyTorch Lightning as already done in Tutorial 5 and 6."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda:0\n"
     ]
    }
   ],
   "source": [
    "## Standard libraries\n",
    "import os\n",
    "import json\n",
    "import math\n",
    "import numpy as np \n",
    "import time\n",
    "\n",
    "## Imports for plotting\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline \n",
    "from IPython.display import set_matplotlib_formats\n",
    "set_matplotlib_formats('svg', 'pdf') # For export\n",
    "from matplotlib.colors import to_rgb\n",
    "import matplotlib\n",
    "matplotlib.rcParams['lines.linewidth'] = 2.0\n",
    "import seaborn as sns\n",
    "sns.reset_orig()\n",
    "sns.set()\n",
    "\n",
    "## Progress bar\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "## PyTorch\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.utils.data as data\n",
    "import torch.optim as optim\n",
    "# Torchvision\n",
    "import torchvision\n",
    "from torchvision.datasets import CIFAR10\n",
    "from torchvision import transforms\n",
    "# PyTorch Lightning\n",
    "try:\n",
    "    import pytorch_lightning as pl\n",
    "except ModuleNotFoundError: # Google Colab does not have PyTorch Lightning installed by default. Hence, we do it here if necessary\n",
    "    !pip install --quiet pytorch-lightning>=1.4\n",
    "    import pytorch_lightning as pl\n",
    "from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint\n",
    "\n",
    "# Path to the folder where the datasets are/should be downloaded (e.g. CIFAR10)\n",
    "DATASET_PATH = \"../data\"\n",
    "# Path to the folder where the pretrained models are saved\n",
    "CHECKPOINT_PATH = \"../saved_models/tutorial7\"\n",
    "\n",
    "# Setting the seed\n",
    "pl.seed_everything(42)\n",
    "\n",
    "# Ensure that all operations are deterministic on GPU (if used) for reproducibility\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False\n",
    "\n",
    "device = torch.device(\"cuda:0\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also have a few pre-trained models we download below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import urllib.request\n",
    "from urllib.error import HTTPError\n",
    "# Github URL where saved models are stored for this tutorial\n",
    "base_url = \"https://raw-hub.myxuebi.top/phlippe/saved_models/main/tutorial7/\"\n",
    "# Files to download\n",
    "pretrained_files = [\"NodeLevelMLP.ckpt\", \"NodeLevelGNN.ckpt\", \"GraphLevelGraphConv.ckpt\"]\n",
    "\n",
    "# Create checkpoint path if it doesn't exist yet\n",
    "os.makedirs(CHECKPOINT_PATH, exist_ok=True)\n",
    "\n",
    "# For each file, check whether it already exists. If not, try downloading it.\n",
    "for file_name in pretrained_files:\n",
    "    file_path = os.path.join(CHECKPOINT_PATH, file_name)\n",
    "    if \"/\" in file_name:\n",
    "        os.makedirs(file_path.rsplit(\"/\",1)[0], exist_ok=True)\n",
    "    if not os.path.isfile(file_path):\n",
    "        file_url = base_url + file_name\n",
    "        print(f\"Downloading {file_url}...\")\n",
    "        try:\n",
    "            urllib.request.urlretrieve(file_url, file_path)\n",
    "        except HTTPError as e:\n",
    "            print(\"Something went wrong. Please try to download the file from the GDrive folder, or contact the author with the full output including the following error:\\n\", e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph Neural Networks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Graph representation\n",
    "\n",
    "Before starting the discussion of specific neural network operations on graphs, we should consider how to represent a graph. Mathematically, a graph $\\mathcal{G}$ is defined as a tuple of a set of nodes/vertices $V$, and a set of edges/links $E$: $\\mathcal{G}=(V,E)$. Each edge is a pair of two vertices, and represents a connection between them. For instance, let's look at the following graph:\n",
    "\n",
    "<center width=\"100%\" style=\"padding:10px\"><img src=\"example_graph.svg\" width=\"250px\"></center>\n",
    "\n",
    "The vertices are $V=\\{1,2,3,4\\}$, and edges $E=\\{(1,2), (2,3), (2,4), (3,4)\\}$. Note that for simplicity, we assume the graph to be undirected and hence don't add mirrored pairs like $(2,1)$. In application, vertices and edge can often have specific attributes, and edges can even be directed. The question is how we could represent this diversity in an efficient way for matrix operations. Usually, for the edges, we decide between two variants: an adjacency matrix, or a list of paired vertex indices. \n",
    "\n",
    "The **adjacency matrix** $A$ is a square matrix whose elements indicate whether pairs of vertices are adjacent, i.e. connected, or not. In the simplest case, $A_{ij}$ is 1 if there is a connection from node $i$ to $j$, and otherwise 0. If we have edge attributes or different categories of edges in a graph, this information can be added to the matrix as well. For an undirected graph, keep in mind that $A$ is a symmetric matrix ($A_{ij}=A_{ji}$). For the example graph above, we have the following adjacency matrix:\n",
    "\n",
    "$$\n",
    "A = \\begin{bmatrix}\n",
    "    0 & 1 & 0 & 0\\\\\n",
    "    1 & 0 & 1 & 1\\\\\n",
    "    0 & 1 & 0 & 1\\\\\n",
    "    0 & 1 & 1 & 0\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "While expressing a graph as a list of edges is more efficient in terms of memory and (possibly) computation, using an adjacency matrix is more intuitive and simpler to implement. In our implementations below, we will rely on the adjacency matrix to keep the code simple. However, common libraries use edge lists, which we will discuss later more.\n",
    "Alternatively, we could also use the list of edges to define a sparse adjacency matrix with which we can work as if it was a dense matrix, but allows more memory-efficient operations. PyTorch supports this with the sub-package `torch.sparse` ([documentation](https://pytorch.org/docs/stable/sparse.html)) which is however still in a beta-stage (API might change in future)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Graph Convolutions\n",
    "\n",
    "Graph Convolutional Networks have been introduced by [Kipf et al.](https://openreview.net/pdf?id=SJU4ayYgl) in 2016 at the University of Amsterdam. He also wrote a great [blog post](https://tkipf.github.io/graph-convolutional-networks/) about this topic, which is recommended if you want to read about GCNs from a different perspective. GCNs are similar to convolutions in images in the sense that the \"filter\" parameters are typically shared over all locations in the graph. At the same time, GCNs rely on message passing methods, which means that vertices exchange information with the neighbors, and send \"messages\" to each other. Before looking at the math, we can try to visually understand how GCNs work. The first step is that each node creates a feature vector that represents the message it wants to send to all its neighbors. In the second step, the messages are sent to the neighbors, so that a node receives one message per adjacent node. Below we have visualized the two steps for our example graph. \n",
    "\n",
    "<center width=\"100%\" style=\"padding:10px\"><img src=\"graph_message_passing.svg\" width=\"700px\"></center>\n",
    "\n",
    "If we want to formulate that in more mathematical terms, we need to first decide how to combine all the messages a node receives. As the number of messages vary across nodes, we need an operation that works for any number. Hence, the usual way to go is to sum or take the mean. Given the previous features of nodes $H^{(l)}$, the GCN layer is defined as follows:\n",
    "\n",
    "$$H^{(l+1)} = \\sigma\\left(\\hat{D}^{-1/2}\\hat{A}\\hat{D}^{-1/2}H^{(l)}W^{(l)}\\right)$$\n",
    "\n",
    "$W^{(l)}$ is the weight parameters with which we transform the input features into messages ($H^{(l)}W^{(l)}$). To the adjacency matrix $A$ we add the identity matrix so that each node sends its own message also to itself: $\\hat{A}=A+I$. Finally, to take the average instead of summing, we calculate the matrix $\\hat{D}$ which is a diagonal matrix with $D_{ii}$ denoting the number of neighbors node $i$ has. $\\sigma$ represents an arbitrary activation function, and not necessarily the sigmoid (usually a ReLU-based activation function is used in GNNs). \n",
    "\n",
    "When implementing the GCN layer in PyTorch, we can take advantage of the flexible operations on tensors. Instead of defining a matrix $\\hat{D}$, we can simply divide the summed messages by the number of neighbors afterward. Additionally, we replace the weight matrix with a linear layer, which additionally allows us to add a bias. Written as a PyTorch module, the GCN layer is defined as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GCNLayer(nn.Module):\n",
    "    \n",
    "    def __init__(self, c_in, c_out):\n",
    "        super().__init__()\n",
    "        self.projection = nn.Linear(c_in, c_out)\n",
    "\n",
    "    def forward(self, node_feats, adj_matrix):\n",
    "        \"\"\"\n",
    "        Inputs:\n",
    "            node_feats - Tensor with node features of shape [batch_size, num_nodes, c_in]\n",
    "            adj_matrix - Batch of adjacency matrices of the graph. If there is an edge from i to j, adj_matrix[b,i,j]=1 else 0.\n",
    "                         Supports directed edges by non-symmetric matrices. Assumes to already have added the identity connections. \n",
    "                         Shape: [batch_size, num_nodes, num_nodes]\n",
    "        \"\"\"\n",
    "        # Num neighbours = number of incoming edges\n",
    "        num_neighbours = adj_matrix.sum(dim=-1, keepdims=True)\n",
    "        node_feats = self.projection(node_feats)\n",
    "        node_feats = torch.bmm(adj_matrix, node_feats)\n",
    "        node_feats = node_feats / num_neighbours\n",
    "        return node_feats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To further understand the GCN layer, we can apply it to our example graph above. First, let's specify some node features and the adjacency matrix with added self-connections:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Node features:\n",
      " tensor([[[0., 1.],\n",
      "         [2., 3.],\n",
      "         [4., 5.],\n",
      "         [6., 7.]]])\n",
      "\n",
      "Adjacency matrix:\n",
      " tensor([[[1., 1., 0., 0.],\n",
      "         [1., 1., 1., 1.],\n",
      "         [0., 1., 1., 1.],\n",
      "         [0., 1., 1., 1.]]])\n"
     ]
    }
   ],
   "source": [
    "node_feats = torch.arange(8, dtype=torch.float32).view(1, 4, 2)\n",
    "adj_matrix = torch.Tensor([[[1, 1, 0, 0],\n",
    "                            [1, 1, 1, 1],\n",
    "                            [0, 1, 1, 1],\n",
    "                            [0, 1, 1, 1]]])\n",
    "\n",
    "print(\"Node features:\\n\", node_feats)\n",
    "print(\"\\nAdjacency matrix:\\n\", adj_matrix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's apply a GCN layer to it. For simplicity, we initialize the linear weight matrix as an identity matrix so that the input features are equal to the messages. This makes it easier for us to verify the message passing operation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adjacency matrix tensor([[[1., 1., 0., 0.],\n",
      "         [1., 1., 1., 1.],\n",
      "         [0., 1., 1., 1.],\n",
      "         [0., 1., 1., 1.]]])\n",
      "Input features tensor([[[0., 1.],\n",
      "         [2., 3.],\n",
      "         [4., 5.],\n",
      "         [6., 7.]]])\n",
      "Output features tensor([[[1., 2.],\n",
      "         [3., 4.],\n",
      "         [4., 5.],\n",
      "         [4., 5.]]])\n"
     ]
    }
   ],
   "source": [
    "layer = GCNLayer(c_in=2, c_out=2)\n",
    "layer.projection.weight.data = torch.Tensor([[1., 0.], [0., 1.]])\n",
    "layer.projection.bias.data = torch.Tensor([0., 0.])\n",
    "\n",
    "with torch.no_grad():\n",
    "    out_feats = layer(node_feats, adj_matrix)\n",
    "\n",
    "print(\"Adjacency matrix\", adj_matrix)\n",
    "print(\"Input features\", node_feats)\n",
    "print(\"Output features\", out_feats)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, the first node's output values are the average of itself and the second node. Similarly, we can verify all other nodes. However, in a GNN, we would also want to allow feature exchange between nodes beyond its neighbors. This can be achieved by applying multiple GCN layers, which gives us the final layout of a GNN. The GNN can be build up by a sequence of GCN layers and non-linearities such as ReLU. For a visualization, see below (figure credit - [Thomas Kipf, 2016](https://tkipf.github.io/graph-convolutional-networks/)).\n",
    "\n",
    "<center width=\"100%\" style=\"padding: 10px\"><img src=\"gcn_network.png\" width=\"600px\"></center>\n",
    "\n",
    "However, one issue we can see from looking at the example above is that the output features for nodes 3 and 4 are the same because they have the same adjacent nodes (including itself). Therefore, GCN layers can make the network forget node-specific information if we just take a mean over all messages. Multiple possible improvements have been proposed. While the simplest option might be using residual connections, the more common approach is to either weigh the self-connections higher or define a separate weight matrix for the self-connections. Alternatively, we can re-visit a concept from the last tutorial: attention. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Graph Attention \n",
    "\n",
    "If you remember from the last tutorial, attention describes a weighted average of multiple elements with the weights dynamically computed based on an input query and elements' keys (if you haven't read Tutorial 6 yet, it is recommended to at least go through the very first section called [What is Attention?](https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/tutorial6/Transformers_and_MHAttention.html#What-is-Attention?)). This concept can be similarly applied to graphs, one of such is the Graph Attention Network (called GAT, proposed by [Velickovic et al., 2017](https://arxiv.org/abs/1710.10903)). Similarly to the GCN, the graph attention layer creates a message for each node using a linear layer/weight matrix. For the attention part, it uses the message from the node itself as a query, and the messages to average as both keys and values (note that this also includes the message to itself). The score function $f_{attn}$ is implemented as a one-layer MLP which maps the query and key to a single value. The MLP looks as follows (figure credit - [Velickovic et al.](https://arxiv.org/abs/1710.10903)):\n",
    "\n",
    "<center width=\"100%\" style=\"padding:10px\"><img src=\"graph_attention_MLP.svg\" width=\"250px\"></center>\n",
    "\n",
    "$h_i$ and $h_j$ are the original features from node $i$ and $j$ respectively, and represent the messages of the layer with $\\mathbf{W}$ as weight matrix. $\\mathbf{a}$ is the weight matrix of the MLP, which has the shape $[1,2\\times d_{\\text{message}}]$, and $\\alpha_{ij}$ the final attention weight from node $i$ to $j$. The calculation can be described as follows:\n",
    "\n",
    "$$\\alpha_{ij} = \\frac{\\exp\\left(\\text{LeakyReLU}\\left(\\mathbf{a}\\left[\\mathbf{W}h_i||\\mathbf{W}h_j\\right]\\right)\\right)}{\\sum_{k\\in\\mathcal{N}_i} \\exp\\left(\\text{LeakyReLU}\\left(\\mathbf{a}\\left[\\mathbf{W}h_i||\\mathbf{W}h_k\\right]\\right)\\right)}$$\n",
    "\n",
    "The operator $||$ represents the concatenation, and $\\mathcal{N}_i$ the indices of the neighbors of node $i$. Note that in contrast to usual practice, we apply a non-linearity (here LeakyReLU) before the softmax over elements. Although it seems like a minor change at first, it is crucial for the attention to depend on the original input. Specifically, let's remove the non-linearity for a second, and try to simplify the expression:\n",
    "\n",
    "$$\n",
    "\\begin{split}\n",
    "    \\alpha_{ij} & = \\frac{\\exp\\left(\\mathbf{a}\\left[\\mathbf{W}h_i||\\mathbf{W}h_j\\right]\\right)}{\\sum_{k\\in\\mathcal{N}_i} \\exp\\left(\\mathbf{a}\\left[\\mathbf{W}h_i||\\mathbf{W}h_k\\right]\\right)}\\\\[5pt]\n",
    "    & = \\frac{\\exp\\left(\\mathbf{a}_{:,:d/2}\\mathbf{W}h_i+\\mathbf{a}_{:,d/2:}\\mathbf{W}h_j\\right)}{\\sum_{k\\in\\mathcal{N}_i} \\exp\\left(\\mathbf{a}_{:,:d/2}\\mathbf{W}h_i+\\mathbf{a}_{:,d/2:}\\mathbf{W}h_k\\right)}\\\\[5pt]\n",
    "    & = \\frac{\\exp\\left(\\mathbf{a}_{:,:d/2}\\mathbf{W}h_i\\right)\\cdot\\exp\\left(\\mathbf{a}_{:,d/2:}\\mathbf{W}h_j\\right)}{\\sum_{k\\in\\mathcal{N}_i} \\exp\\left(\\mathbf{a}_{:,:d/2}\\mathbf{W}h_i\\right)\\cdot\\exp\\left(\\mathbf{a}_{:,d/2:}\\mathbf{W}h_k\\right)}\\\\[5pt]\n",
    "    & = \\frac{\\exp\\left(\\mathbf{a}_{:,d/2:}\\mathbf{W}h_j\\right)}{\\sum_{k\\in\\mathcal{N}_i} \\exp\\left(\\mathbf{a}_{:,d/2:}\\mathbf{W}h_k\\right)}\\\\\n",
    "\\end{split}\n",
    "$$\n",
    "\n",
    "We can see that without the non-linearity, the attention term with $h_i$ actually cancels itself out, resulting in the attention being independent of the node itself. Hence, we would have the same issue as the GCN of creating the same output features for nodes with the same neighbors. This is why the LeakyReLU is crucial and adds some dependency on $h_i$ to the attention. \n",
    "\n",
    "Once we obtain all attention factors, we can calculate the output features for each node by performing the weighted average:\n",
    "\n",
    "$$h_i'=\\sigma\\left(\\sum_{j\\in\\mathcal{N}_i}\\alpha_{ij}\\mathbf{W}h_j\\right)$$\n",
    "\n",
    "$\\sigma$ is yet another non-linearity, as in the GCN layer. Visually, we can represent the full message passing in an attention layer as follows (figure credit - [Velickovic et al.](https://arxiv.org/abs/1710.10903)):\n",
    "\n",
    "<center width=\"100%\"><img src=\"graph_attention.jpeg\" width=\"400px\"></center>\n",
    "\n",
    "To increase the expressiveness of the graph attention network, [Velickovic et al.](https://arxiv.org/abs/1710.10903) proposed to extend it to multiple heads similar to the Multi-Head Attention block in Transformers. This results in $N$ attention layers being applied in parallel. In the image above, it is visualized as three different colors of arrows (green, blue, and purple) that are afterward concatenated. The average is only applied for the very final prediction layer in a network. \n",
    "\n",
    "After having discussed the graph attention layer in detail, we can implement it below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GATLayer(nn.Module):\n",
    "    \n",
    "    def __init__(self, c_in, c_out, num_heads=1, concat_heads=True, alpha=0.2):\n",
    "        \"\"\"\n",
    "        Inputs:\n",
    "            c_in - Dimensionality of input features\n",
    "            c_out - Dimensionality of output features\n",
    "            num_heads - Number of heads, i.e. attention mechanisms to apply in parallel. The \n",
    "                        output features are equally split up over the heads if concat_heads=True.\n",
    "            concat_heads - If True, the output of the different heads is concatenated instead of averaged.\n",
    "            alpha - Negative slope of the LeakyReLU activation.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.num_heads = num_heads\n",
    "        self.concat_heads = concat_heads\n",
    "        if self.concat_heads:\n",
    "            assert c_out % num_heads == 0, \"Number of output features must be a multiple of the count of heads.\"\n",
    "            c_out = c_out // num_heads\n",
    "        \n",
    "        # Sub-modules and parameters needed in the layer\n",
    "        self.projection = nn.Linear(c_in, c_out * num_heads)\n",
    "        self.a = nn.Parameter(torch.Tensor(num_heads, 2 * c_out)) # One per head\n",
    "        self.leakyrelu = nn.LeakyReLU(alpha)\n",
    "        \n",
    "        # Initialization from the original implementation\n",
    "        nn.init.xavier_uniform_(self.projection.weight.data, gain=1.414)\n",
    "        nn.init.xavier_uniform_(self.a.data, gain=1.414)\n",
    "        \n",
    "    def forward(self, node_feats, adj_matrix, print_attn_probs=False):\n",
    "        \"\"\"\n",
    "        Inputs:\n",
    "            node_feats - Input features of the node. Shape: [batch_size, c_in]\n",
    "            adj_matrix - Adjacency matrix including self-connections. Shape: [batch_size, num_nodes, num_nodes]\n",
    "            print_attn_probs - If True, the attention weights are printed during the forward pass (for debugging purposes)\n",
    "        \"\"\"\n",
    "        batch_size, num_nodes = node_feats.size(0), node_feats.size(1)\n",
    "        \n",
    "        # Apply linear layer and sort nodes by head\n",
    "        node_feats = self.projection(node_feats)\n",
    "        node_feats = node_feats.view(batch_size, num_nodes, self.num_heads, -1)\n",
    "        \n",
    "        # We need to calculate the attention logits for every edge in the adjacency matrix \n",
    "        # Doing this on all possible combinations of nodes is very expensive\n",
    "        # => Create a tensor of [W*h_i||W*h_j] with i and j being the indices of all edges\n",
    "        edges = adj_matrix.nonzero(as_tuple=False) # Returns indices where the adjacency matrix is not 0 => edges\n",
    "        node_feats_flat = node_feats.view(batch_size * num_nodes, self.num_heads, -1)\n",
    "        edge_indices_row = edges[:,0] * num_nodes + edges[:,1]\n",
    "        edge_indices_col = edges[:,0] * num_nodes + edges[:,2]\n",
    "        a_input = torch.cat([\n",
    "            torch.index_select(input=node_feats_flat, index=edge_indices_row, dim=0),\n",
    "            torch.index_select(input=node_feats_flat, index=edge_indices_col, dim=0)\n",
    "        ], dim=-1) # Index select returns a tensor with node_feats_flat being indexed at the desired positions along dim=0\n",
    "        \n",
    "        # Calculate attention MLP output (independent for each head)\n",
    "        attn_logits = torch.einsum('bhc,hc->bh', a_input, self.a) \n",
    "        attn_logits = self.leakyrelu(attn_logits)\n",
    "        \n",
    "        # Map list of attention values back into a matrix\n",
    "        attn_matrix = attn_logits.new_zeros(adj_matrix.shape+(self.num_heads,)).fill_(-9e15)\n",
    "        attn_matrix[adj_matrix[...,None].repeat(1,1,1,self.num_heads) == 1] = attn_logits.reshape(-1)\n",
    "        \n",
    "        # Weighted average of attention\n",
    "        attn_probs = F.softmax(attn_matrix, dim=2)\n",
    "        if print_attn_probs:\n",
    "            print(\"Attention probs\\n\", attn_probs.permute(0, 3, 1, 2))\n",
    "        node_feats = torch.einsum('bijh,bjhc->bihc', attn_probs, node_feats)\n",
    "        \n",
    "        # If heads should be concatenated, we can do this by reshaping. Otherwise, take mean\n",
    "        if self.concat_heads:\n",
    "            node_feats = node_feats.reshape(batch_size, num_nodes, -1)\n",
    "        else:\n",
    "            node_feats = node_feats.mean(dim=2)\n",
    "        \n",
    "        return node_feats "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, we can apply the graph attention layer on our example graph above to understand the dynamics better. As before, the input layer is initialized as an identity matrix, but we set $\\mathbf{a}$ to be a vector of arbitrary numbers to obtain different attention values. We use two heads to show the parallel, independent attention mechanisms working in the layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Attention probs\n",
      " tensor([[[[0.3543, 0.6457, 0.0000, 0.0000],\n",
      "          [0.1096, 0.1450, 0.2642, 0.4813],\n",
      "          [0.0000, 0.1858, 0.2885, 0.5257],\n",
      "          [0.0000, 0.2391, 0.2696, 0.4913]],\n",
      "\n",
      "         [[0.5100, 0.4900, 0.0000, 0.0000],\n",
      "          [0.2975, 0.2436, 0.2340, 0.2249],\n",
      "          [0.0000, 0.3838, 0.3142, 0.3019],\n",
      "          [0.0000, 0.4018, 0.3289, 0.2693]]]])\n",
      "Adjacency matrix tensor([[[1., 1., 0., 0.],\n",
      "         [1., 1., 1., 1.],\n",
      "         [0., 1., 1., 1.],\n",
      "         [0., 1., 1., 1.]]])\n",
      "Input features tensor([[[0., 1.],\n",
      "         [2., 3.],\n",
      "         [4., 5.],\n",
      "         [6., 7.]]])\n",
      "Output features tensor([[[1.2913, 1.9800],\n",
      "         [4.2344, 3.7725],\n",
      "         [4.6798, 4.8362],\n",
      "         [4.5043, 4.7351]]])\n"
     ]
    }
   ],
   "source": [
    "layer = GATLayer(2, 2, num_heads=2)\n",
    "layer.projection.weight.data = torch.Tensor([[1., 0.], [0., 1.]])\n",
    "layer.projection.bias.data = torch.Tensor([0., 0.])\n",
    "layer.a.data = torch.Tensor([[-0.2, 0.3], [0.1, -0.1]])\n",
    "\n",
    "with torch.no_grad():\n",
    "    out_feats = layer(node_feats, adj_matrix, print_attn_probs=True)\n",
    "\n",
    "print(\"Adjacency matrix\", adj_matrix)\n",
    "print(\"Input features\", node_feats)\n",
    "print(\"Output features\", out_feats)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We recommend that you try to calculate the attention matrix at least for one head and one node for yourself. The entries are 0 where there does not exist an edge between $i$ and $j$. For the others, we see a diverse set of attention probabilities. Moreover, the output features of node 3 and 4 are now different although they have the same neighbors."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PyTorch Geometric\n",
    "\n",
    "We had mentioned before that implementing graph networks with adjacency matrix is simple and straight-forward but can be computationally expensive for large graphs. Many real-world graphs can reach over 200k nodes, for which adjacency matrix-based implementations fail. There are a lot of optimizations possible when implementing GNNs, and luckily, there exist packages that provide such layers. The most popular packages for PyTorch are [PyTorch Geometric](https://pytorch-geometric.readthedocs.io/en/latest/) and the [Deep Graph Library](https://www.dgl.ai/) (the latter being actually framework agnostic). Which one to use depends on the project you are planning to do and personal taste. In this tutorial, we will look at PyTorch Geometric as part of the PyTorch family. Similar to PyTorch Lightning, PyTorch Geometric is not installed by default on GoogleColab (and actually also not in our `dl2021` environment due to many dependencies that would be unnecessary for the practicals). Hence, let's import and/or install it below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "RDKit WARNING: [19:12:50] Enabling RDKit 2019.09.3 jupyter extensions\n"
     ]
    }
   ],
   "source": [
    "# torch geometric\n",
    "try: \n",
    "    import torch_geometric\n",
    "except ModuleNotFoundError:\n",
    "    # Installing torch geometric packages with specific CUDA+PyTorch version. \n",
    "    # See https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html for details \n",
    "    TORCH = torch.__version__.split('+')[0]\n",
    "    CUDA = 'cu' + torch.version.cuda.replace('.','')\n",
    "\n",
    "    !pip install torch-scatter     -f https://pytorch-geometric.com/whl/torch-{TORCH}+{CUDA}.html\n",
    "    !pip install torch-sparse      -f https://pytorch-geometric.com/whl/torch-{TORCH}+{CUDA}.html\n",
    "    !pip install torch-cluster     -f https://pytorch-geometric.com/whl/torch-{TORCH}+{CUDA}.html\n",
    "    !pip install torch-spline-conv -f https://pytorch-geometric.com/whl/torch-{TORCH}+{CUDA}.html\n",
    "    !pip install torch-geometric \n",
    "    import torch_geometric\n",
    "import torch_geometric.nn as geom_nn\n",
    "import torch_geometric.data as geom_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PyTorch Geometric provides us a set of common graph layers, including the GCN and GAT layer we implemented above. Additionally, similar to PyTorch's torchvision, it provides the common graph datasets and transformations on those to simplify training. Compared to our implementation above, PyTorch Geometric uses a list of index pairs to represent the edges. The details of this library will be explored further in our experiments.\n",
    "\n",
    "In our tasks below, we want to allow us to pick from a multitude of graph layers. Thus, we define again below a dictionary to access those using a string:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "gnn_layer_by_name = {\n",
    "    \"GCN\": geom_nn.GCNConv,\n",
    "    \"GAT\": geom_nn.GATConv,\n",
    "    \"GraphConv\": geom_nn.GraphConv\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additionally to GCN and GAT, we added the layer `geom_nn.GraphConv` ([documentation](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.conv.GraphConv)). GraphConv is a GCN with a separate weight matrix for the self-connections. Mathematically, this would be:\n",
    "\n",
    "$$\n",
    "\\mathbf{x}_i^{(l+1)} = \\mathbf{W}^{(l + 1)}_1 \\mathbf{x}_i^{(l)} + \\mathbf{W}^{(\\ell + 1)}_2 \\sum_{j \\in \\mathcal{N}_i} \\mathbf{x}_j^{(l)}\n",
    "$$\n",
    "\n",
    "In this formula, the neighbor's messages are added instead of averaged. However, PyTorch Geometric provides the argument `aggr` to switch between summing, averaging, and max pooling."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Experiments on graph structures\n",
    "\n",
    "Tasks on graph-structured data can be grouped into three groups: node-level, edge-level and graph-level. The different levels describe on which level we want to perform classification/regression. We will discuss all three types in more detail below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Node-level tasks: Semi-supervised node classification\n",
    "\n",
    "Node-level tasks have the goal to classify nodes in a graph. Usually, we have given a single, large graph with >1000 nodes of which a certain amount of nodes are labeled. We learn to classify those labeled examples during training and try to generalize to the unlabeled nodes. \n",
    "\n",
    "A popular example that we will use in this tutorial is the Cora dataset, a citation network among papers. The Cora consists of 2708 scientific publications with links between each other representing the citation of one paper by another. The task is to classify each publication into one of seven classes. Each publication is represented by a bag-of-words vector. This means that we have a vector of 1433 elements for each publication, where a 1 at feature $i$ indicates that the $i$-th word of a pre-defined dictionary is in the article. Binary bag-of-words representations are commonly used when we need very simple encodings, and already have an intuition of what words to expect in a network. There exist much better approaches, but we will leave this to the NLP courses to discuss.\n",
    "\n",
    "We will load the dataset below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "cora_dataset = torch_geometric.datasets.Planetoid(root=DATASET_PATH, name=\"Cora\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at how PyTorch Geometric represents the graph data. Note that although we have a single graph, PyTorch Geometric returns a dataset for compatibility to other datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(edge_index=[2, 10556], test_mask=[2708], train_mask=[2708], val_mask=[2708], x=[2708, 1433], y=[2708])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cora_dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The graph is represented by a `Data` object ([documentation](https://pytorch-geometric.readthedocs.io/en/latest/modules/data.html#torch_geometric.data.Data)) which we can access as a standard Python namespace. The edge index tensor is the list of edges in the graph and contains the mirrored version of each edge for undirected graphs. The `train_mask`, `val_mask`, and `test_mask` are boolean masks that indicate which nodes we should use for training, validation, and testing. The `x` tensor is the feature tensor of our 2708 publications, and `y` the labels for all nodes.\n",
    "\n",
    "After having seen the data, we can implement a simple graph neural network. The GNN applies a sequence of graph layers (GCN, GAT, or GraphConv), ReLU as activation function, and dropout for regularization. See below for the specific implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GNNModel(nn.Module):\n",
    "    \n",
    "    def __init__(self, c_in, c_hidden, c_out, num_layers=2, layer_name=\"GCN\", dp_rate=0.1, **kwargs):\n",
    "        \"\"\"\n",
    "        Inputs:\n",
    "            c_in - Dimension of input features\n",
    "            c_hidden - Dimension of hidden features\n",
    "            c_out - Dimension of the output features. Usually number of classes in classification\n",
    "            num_layers - Number of \"hidden\" graph layers\n",
    "            layer_name - String of the graph layer to use\n",
    "            dp_rate - Dropout rate to apply throughout the network\n",
    "            kwargs - Additional arguments for the graph layer (e.g. number of heads for GAT)\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        gnn_layer = gnn_layer_by_name[layer_name]\n",
    "        \n",
    "        layers = []\n",
    "        in_channels, out_channels = c_in, c_hidden\n",
    "        for l_idx in range(num_layers-1):\n",
    "            layers += [\n",
    "                gnn_layer(in_channels=in_channels, \n",
    "                          out_channels=out_channels,\n",
    "                          **kwargs),\n",
    "                nn.ReLU(inplace=True),\n",
    "                nn.Dropout(dp_rate)\n",
    "            ]\n",
    "            in_channels = c_hidden\n",
    "        layers += [gnn_layer(in_channels=in_channels, \n",
    "                             out_channels=c_out,\n",
    "                             **kwargs)]\n",
    "        self.layers = nn.ModuleList(layers)\n",
    "    \n",
    "    def forward(self, x, edge_index):\n",
    "        \"\"\"\n",
    "        Inputs:\n",
    "            x - Input features per node\n",
    "            edge_index - List of vertex index pairs representing the edges in the graph (PyTorch geometric notation)\n",
    "        \"\"\"\n",
    "        for l in self.layers:\n",
    "            # For graph layers, we need to add the \"edge_index\" tensor as additional input\n",
    "            # All PyTorch Geometric graph layer inherit the class \"MessagePassing\", hence\n",
    "            # we can simply check the class type.\n",
    "            if isinstance(l, geom_nn.MessagePassing):\n",
    "                x = l(x, edge_index)\n",
    "            else:\n",
    "                x = l(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Good practice in node-level tasks is to create an MLP baseline that is applied to each node independently. This way we can verify whether adding the graph information to the model indeed improves the prediction, or not. It might also be that the features per node are already expressive enough to clearly point towards a specific class. To check this, we implement a simple MLP below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLPModel(nn.Module):\n",
    "    \n",
    "    def __init__(self, c_in, c_hidden, c_out, num_layers=2, dp_rate=0.1):\n",
    "        \"\"\"\n",
    "        Inputs:\n",
    "            c_in - Dimension of input features\n",
    "            c_hidden - Dimension of hidden features\n",
    "            c_out - Dimension of the output features. Usually number of classes in classification\n",
    "            num_layers - Number of hidden layers\n",
    "            dp_rate - Dropout rate to apply throughout the network\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        layers = []\n",
    "        in_channels, out_channels = c_in, c_hidden\n",
    "        for l_idx in range(num_layers-1):\n",
    "            layers += [\n",
    "                nn.Linear(in_channels, out_channels),\n",
    "                nn.ReLU(inplace=True),\n",
    "                nn.Dropout(dp_rate)\n",
    "            ]\n",
    "            in_channels = c_hidden\n",
    "        layers += [nn.Linear(in_channels, c_out)]\n",
    "        self.layers = nn.Sequential(*layers)\n",
    "    \n",
    "    def forward(self, x, *args, **kwargs):\n",
    "        \"\"\"\n",
    "        Inputs:\n",
    "            x - Input features per node\n",
    "        \"\"\"\n",
    "        return self.layers(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can merge the models into a PyTorch Lightning module which handles the training, validation, and testing for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NodeLevelGNN(pl.LightningModule):\n",
    "    \n",
    "    def __init__(self, model_name, **model_kwargs):\n",
    "        super().__init__()\n",
    "        # Saving hyperparameters\n",
    "        self.save_hyperparameters()\n",
    "        \n",
    "        if model_name == \"MLP\":\n",
    "            self.model = MLPModel(**model_kwargs)\n",
    "        else:\n",
    "            self.model = GNNModel(**model_kwargs)\n",
    "        self.loss_module = nn.CrossEntropyLoss()\n",
    "\n",
    "    def forward(self, data, mode=\"train\"):\n",
    "        x, edge_index = data.x, data.edge_index\n",
    "        x = self.model(x, edge_index)\n",
    "        \n",
    "        # Only calculate the loss on the nodes corresponding to the mask\n",
    "        if mode == \"train\":\n",
    "            mask = data.train_mask\n",
    "        elif mode == \"val\":\n",
    "            mask = data.val_mask\n",
    "        elif mode == \"test\":\n",
    "            mask = data.test_mask\n",
    "        else:\n",
    "            assert False, f\"Unknown forward mode: {mode}\"\n",
    "        \n",
    "        loss = self.loss_module(x[mask], data.y[mask])\n",
    "        acc = (x[mask].argmax(dim=-1) == data.y[mask]).sum().float() / mask.sum()\n",
    "        return loss, acc\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        # We use SGD here, but Adam works as well \n",
    "        optimizer = optim.SGD(self.parameters(), lr=0.1, momentum=0.9, weight_decay=2e-3)\n",
    "        return optimizer\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        loss, acc = self.forward(batch, mode=\"train\")\n",
    "        self.log('train_loss', loss)\n",
    "        self.log('train_acc', acc)\n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        _, acc = self.forward(batch, mode=\"val\")\n",
    "        self.log('val_acc', acc)\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        _, acc = self.forward(batch, mode=\"test\")\n",
    "        self.log('test_acc', acc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additionally to the Lightning module, we define a training function below. As we have a single graph, we use a batch size of 1 for the data loader and share the same data loader for the train, validation, and test set (the mask is picked inside the Lightning module). Besides, we set the argument `enable_progress_bar` to False as it usually shows the progress per epoch, but an epoch only consists of a single step. The rest of the code is very similar to what we have seen in Tutorial 5 and 6 already."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_node_classifier(model_name, dataset, **model_kwargs):\n",
    "    pl.seed_everything(42)\n",
    "    node_data_loader = geom_data.DataLoader(dataset, batch_size=1)\n",
    "    \n",
    "    # Create a PyTorch Lightning trainer with the generation callback\n",
    "    root_dir = os.path.join(CHECKPOINT_PATH, \"NodeLevel\" + model_name)\n",
    "    os.makedirs(root_dir, exist_ok=True)\n",
    "    trainer = pl.Trainer(default_root_dir=root_dir,\n",
    "                         callbacks=[ModelCheckpoint(save_weights_only=True, mode=\"max\", monitor=\"val_acc\")],\n",
    "                         accelerator=\"gpu\" if str(device).startswith(\"cuda\") else \"cpu\",\n",
    "                         devices=1,\n",
    "                         max_epochs=200,\n",
    "                         enable_progress_bar=False) # False because epoch size is 1\n",
    "    trainer.logger._default_hp_metric = None # Optional logging argument that we don't need\n",
    "\n",
    "    # Check whether pretrained model exists. If yes, load it and skip training\n",
    "    pretrained_filename = os.path.join(CHECKPOINT_PATH, f\"NodeLevel{model_name}.ckpt\")\n",
    "    if os.path.isfile(pretrained_filename):\n",
    "        print(\"Found pretrained model, loading...\")\n",
    "        model = NodeLevelGNN.load_from_checkpoint(pretrained_filename)\n",
    "    else:\n",
    "        pl.seed_everything()\n",
    "        model = NodeLevelGNN(model_name=model_name, c_in=dataset.num_node_features, c_out=dataset.num_classes, **model_kwargs)\n",
    "        trainer.fit(model, node_data_loader, node_data_loader)\n",
    "        model = NodeLevelGNN.load_from_checkpoint(trainer.checkpoint_callback.best_model_path)\n",
    "    \n",
    "    # Test best model on the test set\n",
    "    test_result = trainer.test(model, node_data_loader, verbose=False)\n",
    "    batch = next(iter(node_data_loader))\n",
    "    batch = batch.to(model.device)\n",
    "    _, train_acc = model.forward(batch, mode=\"train\")\n",
    "    _, val_acc = model.forward(batch, mode=\"val\")\n",
    "    result = {\"train\": train_acc,\n",
    "              \"val\": val_acc,\n",
    "              \"test\": test_result[0]['test_acc']}\n",
    "    return model, result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can train our models. First, let's train the simple MLP:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Small function for printing the test scores\n",
    "def print_results(result_dict):\n",
    "    if \"train\" in result_dict:\n",
    "        print(f\"Train accuracy: {(100.0*result_dict['train']):4.2f}%\")\n",
    "    if \"val\" in result_dict:\n",
    "        print(f\"Val accuracy:   {(100.0*result_dict['val']):4.2f}%\")\n",
    "    print(f\"Test accuracy:  {(100.0*result_dict['test']):4.2f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n",
      "WARNING: Logging before flag parsing goes to stderr.\n",
      "I1113 19:12:52.401648 139969460983616 distributed.py:49] GPU available: True, used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "I1113 19:12:52.403518 139969460983616 distributed.py:49] TPU available: False, using: 0 TPU cores\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "I1113 19:12:52.404974 139969460983616 accelerator_connector.py:385] LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found pretrained model, loading...\n",
      "Train accuracy: 96.43%\n",
      "Val accuracy:   52.60%\n",
      "Test accuracy:  60.60%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/phillip/anaconda3/envs/nlp1/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:45: UserWarning: The dataloader, test dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 16 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  warnings.warn(*args, **kwargs)\n"
     ]
    }
   ],
   "source": [
    "node_mlp_model, node_mlp_result = train_node_classifier(model_name=\"MLP\",\n",
    "                                                        dataset=cora_dataset,\n",
    "                                                        c_hidden=16,\n",
    "                                                        num_layers=2,\n",
    "                                                        dp_rate=0.1)\n",
    "\n",
    "print_results(node_mlp_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Although the MLP can overfit on the training dataset because of the high-dimensional input features, it does not perform too well on the test set. Let's see if we can beat this score with our graph networks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n",
      "I1113 19:12:53.906714 139969460983616 distributed.py:49] GPU available: True, used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "I1113 19:12:53.907762 139969460983616 distributed.py:49] TPU available: False, using: 0 TPU cores\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "I1113 19:12:53.909639 139969460983616 accelerator_connector.py:385] LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found pretrained model, loading...\n",
      "Train accuracy: 100.00%\n",
      "Val accuracy:   77.20%\n",
      "Test accuracy:  82.40%\n"
     ]
    }
   ],
   "source": [
    "node_gnn_model, node_gnn_result = train_node_classifier(model_name=\"GNN\",\n",
    "                                                        layer_name=\"GCN\",\n",
    "                                                        dataset=cora_dataset, \n",
    "                                                        c_hidden=16, \n",
    "                                                        num_layers=2,\n",
    "                                                        dp_rate=0.1)\n",
    "print_results(node_gnn_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we would have hoped for, the GNN model outperforms the MLP by quite a margin. This shows that using the graph information indeed improves our predictions and lets us generalizes better.\n",
    "\n",
    "The hyperparameters in the model have been chosen to create a relatively small network. This is because the first layer with an input dimension of 1433 can be relatively expensive to perform for large graphs. In general, GNNs can become relatively expensive for very big graphs. This is why such GNNs either have a small hidden size or use a special batching strategy where we sample a connected subgraph of the big, original graph. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Edge-level tasks: Link prediction\n",
    "\n",
    "In some applications, we might have to predict on an edge-level instead of node-level. The most common edge-level task in GNN is link prediction. Link prediction means that given a graph, we want to predict whether there will be/should be an edge between two nodes or not. For example, in a social network, this is used by Facebook and co to propose new friends to you. Again, graph level information can be crucial to perform this task. The output prediction is usually done by performing a similarity metric on the pair of node features, which should be 1 if there should be a link, and otherwise close to 0. To keep the tutorial short, we will not implement this task ourselves. Nevertheless, there are many good resources out there if you are interested in looking closer at this task.\n",
    "Tutorials and papers for this topic include:\n",
    "\n",
    "* [PyTorch Geometric example](https://github.com/rusty1s/pytorch_geometric/blob/master/examples/link_pred.py)\n",
    "* [Graph Neural Networks: A Review of Methods and Applications](https://arxiv.org/pdf/1812.08434.pdf), Zhou et al. 2019\n",
    "* [Link Prediction Based on Graph Neural Networks](https://papers.nips.cc/paper/2018/file/53f0d7c537d99b3824f0f99d62ea2428-Paper.pdf), Zhang and Chen, 2018."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Graph-level tasks: Graph classification\n",
    "\n",
    "Finally, in this part of the tutorial, we will have a closer look at how to apply GNNs to the task of graph classification. The goal is to classify an entire graph instead of single nodes or edges. Therefore, we are also given a dataset of multiple graphs that we need to classify based on some structural graph properties. The most common task for graph classification is molecular property prediction, in which molecules are represented as graphs. Each atom is linked to a node, and edges in the graph are the bonds between atoms. For example, look at the figure below. \n",
    "\n",
    "<center width=\"100%\"><img src=\"molecule_graph.svg\" width=\"600px\"></center>\n",
    "\n",
    "On the left, we have an arbitrary, small molecule with different atoms, whereas the right part of the image shows the graph representation. The atom types are abstracted as node features (e.g. a one-hot vector), and the different bond types are used as edge features. For simplicity, we will neglect the edge attributes in this tutorial, but you can include by using methods like the [Relational Graph Convolution](https://arxiv.org/abs/1703.06103) that uses a different weight matrix for each edge type.\n",
    "\n",
    "The dataset we will use below is called the MUTAG dataset. It is a common small benchmark for graph classification algorithms, and contain 188 graphs with 18 nodes and 20 edges on average for each graph. The graph nodes have 7 different labels/atom types, and the binary graph labels represent \"their mutagenic effect on a specific gram negative bacterium\" (the specific meaning of the labels are not too important here). The dataset is part of a large collection of different graph classification datasets, known as the [TUDatasets](https://chrsmrrs.github.io/datasets/), which is directly accessible via `torch_geometric.datasets.TUDataset` ([documentation](https://pytorch-geometric.readthedocs.io/en/latest/modules/datasets.html#torch_geometric.datasets.TUDataset)) in PyTorch Geometric. We can load the dataset below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "tu_dataset = torch_geometric.datasets.TUDataset(root=DATASET_PATH, name=\"MUTAG\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at some statistics for the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data object: Data(edge_attr=[7442, 4], edge_index=[2, 7442], x=[3371, 7], y=[188])\n",
      "Length: 188\n",
      "Average label: 0.66\n"
     ]
    }
   ],
   "source": [
    "print(\"Data object:\", tu_dataset.data)\n",
    "print(\"Length:\", len(tu_dataset))\n",
    "print(f\"Average label: {tu_dataset.data.y.float().mean().item():4.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first line shows how the dataset stores different graphs. The nodes, edges, and labels of each graph are concatenated to one tensor, and the dataset stores the indices where to split the tensors correspondingly. The length of the dataset is the number of graphs we have, and the \"average label\" denotes the percentage of the graph with label 1. As long as the percentage is in the range of 0.5, we have a relatively balanced dataset. It happens quite often that graph datasets are very imbalanced, hence checking the class balance is always a good thing to do.\n",
    "\n",
    "Next, we will split our dataset into a training and test part. Note that we do not use a validation set this time because of the small size of the dataset. Therefore, our model might overfit slightly on the validation set due to the noise of the evaluation, but we still get an estimate of the performance on untrained data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(42)\n",
    "tu_dataset.shuffle()\n",
    "train_dataset = tu_dataset[:150]\n",
    "test_dataset = tu_dataset[150:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When using a data loader, we encounter a problem with batching $N$ graphs. Each graph in the batch can have a different number of nodes and edges, and hence we would require a lot of padding to obtain a single tensor. Torch geometric uses a different, more efficient approach: we can view the $N$ graphs in a batch as a single large graph with concatenated node and edge list. As there is no edge between the $N$ graphs, running GNN layers on the large graph gives us the same output as running the GNN on each graph separately. Visually, this batching strategy is visualized below (figure credit - PyTorch Geometric team, [tutorial here](https://colab.research.google.com/drive/1I8a0DfQ3fI7Njc62__mVXUlcAleUclnb?usp=sharing#scrollTo=2owRWKcuoALo)).\n",
    "\n",
    "<center width=\"100%\"><img src=\"torch_geometric_stacking_graphs.png\" width=\"600px\"></center>\n",
    "\n",
    "The adjacency matrix is zero for any nodes that come from two different graphs, and otherwise according to the adjacency matrix of the individual graph. Luckily, this strategy is already implemented in torch geometric, and hence we can use the corresponding data loader:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_train_loader = geom_data.DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
    "graph_val_loader = geom_data.DataLoader(test_dataset, batch_size=64) # Additional loader if you want to change to a larger dataset\n",
    "graph_test_loader = geom_data.DataLoader(test_dataset, batch_size=64)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load a batch below to see the batching in action:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch: Batch(batch=[687], edge_attr=[1512, 4], edge_index=[2, 1512], x=[687, 7], y=[38])\n",
      "Labels: tensor([1, 1, 1, 0, 0, 0, 1, 1, 1, 0])\n",
      "Batch indices: tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
      "        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2])\n"
     ]
    }
   ],
   "source": [
    "batch = next(iter(graph_test_loader))\n",
    "print(\"Batch:\", batch)\n",
    "print(\"Labels:\", batch.y[:10])\n",
    "print(\"Batch indices:\", batch.batch[:40])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have 38 graphs stacked together for the test dataset. The batch indices, stored in `batch`, show that the first 12 nodes belong to the first graph, the next 22 to the second graph, and so on. These indices are important for performing the final prediction. To perform a prediction over a whole graph, we usually perform a pooling operation over all nodes after running the GNN model. In this case, we will use the average pooling. Hence, we need to know which nodes should be included in which average pool. Using this pooling, we can already create our graph network below. Specifically, we re-use our class `GNNModel` from before, and simply add an average pool and single linear layer for the graph prediction task. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GraphGNNModel(nn.Module):\n",
    "    \n",
    "    def __init__(self, c_in, c_hidden, c_out, dp_rate_linear=0.5, **kwargs):\n",
    "        \"\"\"\n",
    "        Inputs:\n",
    "            c_in - Dimension of input features\n",
    "            c_hidden - Dimension of hidden features\n",
    "            c_out - Dimension of output features (usually number of classes)\n",
    "            dp_rate_linear - Dropout rate before the linear layer (usually much higher than inside the GNN)\n",
    "            kwargs - Additional arguments for the GNNModel object\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.GNN = GNNModel(c_in=c_in, \n",
    "                            c_hidden=c_hidden, \n",
    "                            c_out=c_hidden, # Not our prediction output yet!\n",
    "                            **kwargs)\n",
    "        self.head = nn.Sequential(\n",
    "            nn.Dropout(dp_rate_linear),\n",
    "            nn.Linear(c_hidden, c_out)\n",
    "        )\n",
    "\n",
    "    def forward(self, x, edge_index, batch_idx):\n",
    "        \"\"\"\n",
    "        Inputs:\n",
    "            x - Input features per node\n",
    "            edge_index - List of vertex index pairs representing the edges in the graph (PyTorch geometric notation)\n",
    "            batch_idx - Index of batch element for each node\n",
    "        \"\"\"\n",
    "        x = self.GNN(x, edge_index)\n",
    "        x = geom_nn.global_mean_pool(x, batch_idx) # Average pooling\n",
    "        x = self.head(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can create a PyTorch Lightning module to handle the training. It is similar to the modules we have seen before and does nothing surprising in terms of training. As we have a binary classification task, we use the Binary Cross Entropy loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GraphLevelGNN(pl.LightningModule):\n",
    "    \n",
    "    def __init__(self, **model_kwargs):\n",
    "        super().__init__()\n",
    "        # Saving hyperparameters\n",
    "        self.save_hyperparameters()\n",
    "        \n",
    "        self.model = GraphGNNModel(**model_kwargs)\n",
    "        self.loss_module = nn.BCEWithLogitsLoss() if self.hparams.c_out == 1 else nn.CrossEntropyLoss()\n",
    "\n",
    "    def forward(self, data, mode=\"train\"):\n",
    "        x, edge_index, batch_idx = data.x, data.edge_index, data.batch\n",
    "        x = self.model(x, edge_index, batch_idx)\n",
    "        x = x.squeeze(dim=-1)\n",
    "        \n",
    "        if self.hparams.c_out == 1:\n",
    "            preds = (x > 0).float()\n",
    "            data.y = data.y.float()\n",
    "        else:\n",
    "            preds = x.argmax(dim=-1)\n",
    "        loss = self.loss_module(x, data.y)\n",
    "        acc = (preds == data.y).sum().float() / preds.shape[0]\n",
    "        return loss, acc\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = optim.AdamW(self.parameters(), lr=1e-2, weight_decay=0.0) # High lr because of small dataset and small model\n",
    "        return optimizer\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        loss, acc = self.forward(batch, mode=\"train\")\n",
    "        self.log('train_loss', loss)\n",
    "        self.log('train_acc', acc)\n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        _, acc = self.forward(batch, mode=\"val\")\n",
    "        self.log('val_acc', acc)\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        _, acc = self.forward(batch, mode=\"test\")\n",
    "        self.log('test_acc', acc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below we train the model on our dataset. It resembles the typical training functions we have seen so far."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_graph_classifier(model_name, **model_kwargs):\n",
    "    pl.seed_everything(42)\n",
    "    \n",
    "    # Create a PyTorch Lightning trainer with the generation callback\n",
    "    root_dir = os.path.join(CHECKPOINT_PATH, \"GraphLevel\" + model_name)\n",
    "    os.makedirs(root_dir, exist_ok=True)\n",
    "    trainer = pl.Trainer(default_root_dir=root_dir,\n",
    "                         callbacks=[ModelCheckpoint(save_weights_only=True, mode=\"max\", monitor=\"val_acc\")],\n",
    "                         accelerator=\"gpu\" if str(device).startswith(\"cuda\") else \"cpu\",\n",
    "                         devices=1,\n",
    "                         max_epochs=500,\n",
    "                         enable_progress_bar=False)\n",
    "    trainer.logger._default_hp_metric = None # Optional logging argument that we don't need\n",
    "\n",
    "    # Check whether pretrained model exists. If yes, load it and skip training\n",
    "    pretrained_filename = os.path.join(CHECKPOINT_PATH, f\"GraphLevel{model_name}.ckpt\")\n",
    "    if os.path.isfile(pretrained_filename):\n",
    "        print(\"Found pretrained model, loading...\")\n",
    "        model = GraphLevelGNN.load_from_checkpoint(pretrained_filename)\n",
    "    else:\n",
    "        pl.seed_everything(42)\n",
    "        model = GraphLevelGNN(c_in=tu_dataset.num_node_features, \n",
    "                              c_out=1 if tu_dataset.num_classes==2 else tu_dataset.num_classes, \n",
    "                              **model_kwargs)\n",
    "        trainer.fit(model, graph_train_loader, graph_val_loader)\n",
    "        model = GraphLevelGNN.load_from_checkpoint(trainer.checkpoint_callback.best_model_path)\n",
    "    # Test best model on validation and test set\n",
    "    train_result = trainer.test(model, graph_train_loader, verbose=False)\n",
    "    test_result = trainer.test(model, graph_test_loader, verbose=False)\n",
    "    result = {\"test\": test_result[0]['test_acc'], \"train\": train_result[0]['test_acc']} \n",
    "    return model, result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let's perform the training and testing. Feel free to experiment with different GNN layers, hyperparameters, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n",
      "I1113 19:12:54.045717 139969460983616 distributed.py:49] GPU available: True, used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "I1113 19:12:54.047091 139969460983616 distributed.py:49] TPU available: False, using: 0 TPU cores\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "I1113 19:12:54.048336 139969460983616 accelerator_connector.py:385] LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "/home/phillip/anaconda3/envs/nlp1/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:45: UserWarning: Your test_dataloader has `shuffle=True`, it is best practice to turn this off for validation and test dataloaders.\n",
      "  warnings.warn(*args, **kwargs)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found pretrained model, loading...\n"
     ]
    }
   ],
   "source": [
    "model, result = train_graph_classifier(model_name=\"GraphConv\", \n",
    "                                       c_hidden=256, \n",
    "                                       layer_name=\"GraphConv\", \n",
    "                                       num_layers=3, \n",
    "                                       dp_rate_linear=0.5,\n",
    "                                       dp_rate=0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train performance: 94.27%\n",
      "Test performance:  92.11%\n"
     ]
    }
   ],
   "source": [
    "print(f\"Train performance: {100.0*result['train']:4.2f}%\")\n",
    "print(f\"Test performance:  {100.0*result['test']:4.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The test performance shows that we obtain quite good scores on an unseen part of the dataset. It should be noted that as we have been using the test set for validation as well, we might have overfitted slightly to this set. Nevertheless, the experiment shows us that GNNs can be indeed powerful to predict the properties of graphs and/or molecules."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "In this tutorial, we have seen the application of neural networks to graph structures. We looked at how a graph can be represented (adjacency matrix or edge list), and discussed the implementation of common graph layers: GCN and GAT. The implementations showed the practical side of the layers, which is often easier than the theory. Finally, we experimented with different tasks, on node-, edge- and graph-level. Overall, we have seen that including graph information in the predictions can be crucial for achieving high performance. There are a lot of applications that benefit from GNNs, and the importance of these networks will likely increase over the next years."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "[![Star our repository](https://img.shields.io/static/v1.svg?logo=star&label=⭐&message=Star%20Our%20Repository&color=yellow)](https://github.com/phlippe/uvadlc_notebooks/)  If you found this tutorial helpful, consider ⭐-ing our repository.    \n",
    "[![Ask questions](https://img.shields.io/static/v1.svg?logo=star&label=❔&message=Ask%20Questions&color=9cf)](https://github.com/phlippe/uvadlc_notebooks/issues)  For any questions, typos, or bugs that you found, please raise an issue on GitHub. \n",
    "\n",
    "---"
   ]
  }
 ],
 "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}