GNN-tutorial-solution.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Graph Neural Networks Tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1. Pytorch Geometric Framework\n",
    "- Understanding Message Passing Scheme in Pytorch Geometric.\n",
    "- Efficient graph data representations and paralleling minibatching graphs.\n",
    "- Showcase the implementation of **Graph Convolution Networks** (Kipf & Welling, [SEMI-SUPERVISED CLASSIFICATION WITH GRAPH CONVOLUTIONAL NETWORKS](https://arxiv.org/abs/1609.02907), ICLR 2017), and you should implement **GraphSAGE** (Hamilton et al, [Inductive Representation Learning on Large Graphs](https://arxiv.org/abs/1706.02216), NIPS 2017) in the lab based on message passing scheme.\n",
    "\n",
    "#### 2. Vertex Classification\n",
    "- Showcase a model developed based on our GCN implementation to do vertex classification on Cora dataset. \n",
    "- Develop a model with **your own** GraphSAGE (with mean/sum/max aggregation) implementation on the same dataset to get insights of difference.\n",
    "\n",
    "#### 3. Graph Classification\n",
    "- Implement **GINConv** (Xu et al, [HOW POWERFUL ARE GRAPH NEURAL NETWORKS?](https://arxiv.org/abs/1810.00826), ICLR 2019) on graph classification benchmark dataset (IMDB) and compare different aggregation functions (SUM/MEAN/MAX)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up working environment\n",
    "\n",
    "For this tutorial you will need to train a large network, therefore we recommend you work with Google Colaboratory, which provides free GPU time. You will need a Google account to do so. Please log in to your account and go to the following page: https://colab.research.google.com. Then upload this notebook.For GPU support, go to \"Edit\" -> \"Notebook Settings\", and select \"Hardware accelerator\" as \"GPU\".You will need to install pytorch by running the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torch torchvision\n",
    "!pip install torch-scatter torch-sparse torch-geometric"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pytorch Geometric Framework\n",
    "\n",
    "#### Generic Message Passing Scheme\n",
    "Generalizing the convolution operator to irregular domains is typically expressed as a *neighborhood aggregation* or *message passing* scheme.\n",
    "With $\\mathbf{x}^{(k-1)}_i \\in \\mathbb{R}^F$ denoting node features of node $i$ in layer $(k-1)$ and $\\mathbf{e}_{i,j} \\in \\mathbb{R}^D$ denoting (optional) edge features from node $i$ to node $j$, message passing graph neural networks can be described as\n",
    "\n",
    "$$\n",
    "  \\mathbf{x}_i^{(k)} = \\gamma^{(k)} \\left( \\mathbf{x}_i^{(k-1)}, \\square_{j \\in \\mathcal{N}(i)} \\, \\phi^{(k)}\\left(\\mathbf{x}_i^{(k-1)}, \\mathbf{x}_j^{(k-1)},\\mathbf{e}_{i,j}\\right) \\right)\n",
    "$$\n",
    "\n",
    "where $\\square$ denotes a differentiable, permutation invariant function, *e.g.*, sum, mean or max, and $\\gamma$ and $\\phi$ denote differentiable functions such as MLPs (Multi Layer Perceptrons).\n",
    "\n",
    "#### Graph data representations in PyG\n",
    "Given a *sparse* **Graph** $\\mathcal{G}=(\\mathbf{X}, (\\mathbf{I}, \\mathbf{E}))$ with **node features** $\\mathbf{X} \\in \\mathbb{R}^{|V| \\times F}$, **edge indices $\\mathbf{I} \\in \\{1, \\cdots, N\\}^{2 \\times |\\mathcal{E}|}$**, (optional) **edge features** $\\mathbf{E} \\in \\mathbb{R}^{|\\mathcal{E} \\times D|}$, it is described by an instance of class `torch_geometric.data.Data`, which holds the corresponding attributes.\n",
    "\n",
    "We show a simple example of an unweighted and directed graph with four nodes and three edges.\n",
    "\n",
    "<p align=\"center\"><img width=\"70%\" src=\"figures/graph_data.png\"></p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(edge_index=[2, 3], x=[3, 1])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from torch_geometric.data import Data\n",
    "\n",
    "edge_index = torch.tensor([[2, 1, 3],\n",
    "                           [0, 0, 2]], dtype=torch.long)\n",
    "x = torch.tensor([[1], [1], [1]], dtype=torch.float)\n",
    "\n",
    "data = Data(x=x, edge_index=edge_index)\n",
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Mini-Batching Graphs\n",
    "Neural networks are usually trained in a batch-wise fashion. Minibatch graphs can be efficiently dealt with to achieve parallelization over a mini-batch from creating sparse block diagnoal adjacency matrices and concatenating features and target matrices in the node dimension.\n",
    "\n",
    "\n",
    "<p align=\"center\"><img width=\"70%\" src=\"figures/mini_batch_graph.png\"></p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Abstract Message Passing Scheme in PyG\n",
    "\n",
    "PyTorch Geometric provides the `torch_geometric.nn.MessagePassing` base class, which helps in creating such kinds of message passing graph neural networks by automatically taking care of message propagation. The implementation is decoupled into **UPDATE**, **AGGREGATION**, **MESSAGE** functions as:\n",
    "$$\n",
    "    \\mathbf{x}_i^{(k)} = \\mathrm{UPDATE} \\left( \\mathbf{x}_i, , \\mathrm{AGGR}_{j \\in \\mathcal{N}(i)} \\, \\mathrm{MESSAGE}^{(k)}\\left(\\mathbf{x}_i^{(k-1)}, \\mathbf{x}_j^{(k-1)},\\mathbf{e}_{i,j}\\right) \\right)    \n",
    "$$\n",
    "\n",
    "<p align=\"center\"><img width=\"70%\" src=\"figures/message_passing.png\"></p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Implementing the GCN layer (lecture)\n",
    "\n",
    "The graph convolutional operator introduced by Kipf & Welling (ICLR 2017) is defined as\n",
    "$$\n",
    "        \\mathbf{X}^{k} = \\mathbf{\\hat{D}}^{-1/2} \\mathbf{\\hat{A}}\n",
    "        \\mathbf{\\hat{D}}^{-1/2} \\mathbf{X}^{k-1} \\mathbf{\\Theta},\n",
    "$$\n",
    "where $\\mathbf{\\hat{A}} = \\mathbf{A} + \\mathbf{I}$ denotes the adjacency matrix with inserted self-loops and\n",
    "$\\hat{D}_{ii} = \\sum_{j=0} \\hat{A}_{ij}$ its diagonal degree matrix. It is equivalent as:\n",
    "$$\n",
    "\\mathbf{x}_i^{(k)} = \\sum_{j \\in \\mathcal{N}(i) \\cup \\{ i \\}} \\frac{1}{\\sqrt{\\deg(i)} \\cdot \\sqrt{deg(j)}} \\cdot \\left( \\mathbf{x}_j^{(k-1)}\\mathbf{\\Theta} \\right),\n",
    "$$\n",
    "\n",
    "where neighboring node features are first transformed by a weight matrix $\\mathbf{\\Theta}$, normalized by their degree, and finally summed up.\n",
    "This formula can be divided into the following steps:\n",
    "\n",
    "1. Add self-loops to the adjacency matrix.\n",
    "2. Linearly transform node feature matrix.\n",
    "3. Normalize node features.\n",
    "4. Sum up neighboring node features.\n",
    "5. Return new node embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch_geometric.nn import MessagePassing\n",
    "import math\n",
    "\n",
    "def glorot(tensor):\n",
    "    if tensor is not None:\n",
    "        stdv = math.sqrt(6.0 / (tensor.size(-2) + tensor.size(-1)))\n",
    "        tensor.data.uniform_(-stdv, stdv)\n",
    "\n",
    "\n",
    "def zeros(tensor):\n",
    "    if tensor is not None:\n",
    "        tensor.data.fill_(0)\n",
    "\n",
    "        \n",
    "def add_self_loops(edge_index, num_nodes=None):\n",
    "    loop_index = torch.arange(0, num_nodes, dtype=torch.long,\n",
    "                              device=edge_index.device)\n",
    "    loop_index = loop_index.unsqueeze(0).repeat(2, 1)\n",
    "\n",
    "    edge_index = torch.cat([edge_index, loop_index], dim=1)\n",
    "\n",
    "    return edge_index\n",
    "\n",
    "\n",
    "def degree(index, num_nodes=None, dtype=None):\n",
    "    out = torch.zeros((num_nodes), dtype=dtype, device=index.device)\n",
    "    return out.scatter_add_(0, index, out.new_ones((index.size(0))))\n",
    "        \n",
    "\n",
    "class GCNConv(MessagePassing):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super(GCNConv, self).__init__(aggr='add')  # \"Add\" aggregation.\n",
    "        self.lin = torch.nn.Linear(in_channels, out_channels)\n",
    "        \n",
    "        self.reset_parameters()\n",
    "        \n",
    "    def reset_parameters(self):\n",
    "        glorot(self.lin.weight)\n",
    "        zeros(self.lin.bias)\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        # x has shape [N, in_channels]\n",
    "        # edge_index has shape [2, E]\n",
    "        \n",
    "        ########################################################################\n",
    "        #      START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)             #\n",
    "        ########################################################################\n",
    "        # Step 1: Add self-loops to the adjacency matrix.\n",
    "        \n",
    "        edge_index = add_self_loops(edge_index, num_nodes=x.size(0))\n",
    "\n",
    "        # Step 2: Linearly transform node feature matrix.\n",
    "        x = self.lin(x)\n",
    "\n",
    "        # Step 3-5: Start propagating messages.\n",
    "\n",
    "        return self.propagate(edge_index, x=x)\n",
    "        ########################################################################\n",
    "        #                             END OF YOUR CODE                         #\n",
    "        ########################################################################                \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    def message(self, x_j, edge_index, size):\n",
    "        # x_j has shape [E, out_channels]\n",
    "\n",
    "        ########################################################################\n",
    "        #      START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)             #\n",
    "        ########################################################################\n",
    "\n",
    "        # Step 3: Normalize node features.\n",
    "        row, col = edge_index\n",
    "        deg = degree(row, size[0], dtype=x_j.dtype)\n",
    "        deg_inv_sqrt = deg.pow(-0.5)\n",
    "        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0\n",
    "        norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]\n",
    "\n",
    "        return norm.view(-1, 1) * x_j        \n",
    "        \n",
    "        ########################################################################\n",
    "        #                             END OF YOUR CODE                         #\n",
    "        ########################################################################              \n",
    "        \n",
    "\n",
    "\n",
    "    def update(self, aggr_out):\n",
    "        # aggr_out has shape [N, out_channels]\n",
    "\n",
    "        # Step 5: Return new node embeddings.\n",
    "        return aggr_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Implementing GraphSAGE (lab)\n",
    "\n",
    "The algorithm of GraphSAGE (*Inductive Representation Learning on Large Graphs (NIPS 2017)*) embedding generation is described as:\n",
    "\n",
    "<p align=\"center\"><img width=\"70%\" src=\"graphsage.png\"></p>\n",
    "\n",
    "You are required to implement this algortihm with **MEAN/SUM/MAX** AGGREGATE."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.nn import Parameter\n",
    "from torch_geometric.nn.conv import MessagePassing\n",
    "\n",
    "def uniform(size, tensor):\n",
    "    bound = 1.0 / math.sqrt(size)\n",
    "    if tensor is not None:\n",
    "        tensor.data.uniform_(-bound, bound)\n",
    "\n",
    "\n",
    "class SAGEConv(MessagePassing):\n",
    "    def __init__(self, in_channels, out_channels, aggr):\n",
    "        super(SAGEConv, self).__init__(aggr=aggr)\n",
    "\n",
    "        self.in_channels = in_channels\n",
    "        self.out_channels = out_channels\n",
    "\n",
    "        self.weight = Parameter(torch.Tensor(2 * in_channels, out_channels))\n",
    "        \n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        uniform(self.weight.size(0), self.weight)\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        \n",
    "        ########################################################################\n",
    "        #      START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)             #\n",
    "        ########################################################################\n",
    "\n",
    "        return self.propagate(edge_index, x=x)        \n",
    "        \n",
    "        ########################################################################\n",
    "        #                             END OF YOUR CODE                         #\n",
    "        ########################################################################                    \n",
    "\n",
    "\n",
    "\n",
    "    def message(self, x_j, edge_weight):\n",
    "        \n",
    "        ########################################################################\n",
    "        #      START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)             #\n",
    "        ########################################################################\n",
    "\n",
    "        return x_j        \n",
    "        \n",
    "        ########################################################################\n",
    "        #                             END OF YOUR CODE                         #\n",
    "        ########################################################################                    \n",
    "        \n",
    "\n",
    "\n",
    "    def update(self, aggr_out, x):\n",
    "        \n",
    "        \n",
    "        ########################################################################\n",
    "        #      START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)             #\n",
    "        ########################################################################\n",
    "\n",
    "        aggr_out = torch.cat([x, aggr_out], dim=-1)\n",
    "        aggr_out = torch.matmul(aggr_out, self.weight)\n",
    "        aggr_out = F.normalize(aggr_out, p=2, dim=-1)\n",
    "\n",
    "        return aggr_out        \n",
    "        \n",
    "        ########################################################################\n",
    "        #                             END OF YOUR CODE                         #\n",
    "        ########################################################################                    \n",
    "        \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vertex Classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import os.path as osp\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.datasets import Planetoid\n",
    "import torch_geometric.transforms as T\n",
    "\n",
    "path = osp.join(os.getcwd(), 'data', 'Cora')\n",
    "dataset = Planetoid(path, 'Cora')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "from torch import tensor\n",
    "from torch.optim import Adam\n",
    "\n",
    "def run(dataset, model, runs, epochs, lr, weight_decay, early_stopping):\n",
    "\n",
    "    val_losses, accs, durations = [], [], []\n",
    "    for _ in range(runs):\n",
    "        data = dataset[0]\n",
    "        data = data.to(device)\n",
    "\n",
    "        model.to(device).reset_parameters()\n",
    "        optimizer = Adam(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
    "\n",
    "        if torch.cuda.is_available():\n",
    "            torch.cuda.synchronize()\n",
    "\n",
    "        t_start = time.perf_counter()\n",
    "\n",
    "        best_val_loss = float('inf')\n",
    "        test_acc = 0\n",
    "        val_loss_history = []\n",
    "\n",
    "        for epoch in range(1, epochs + 1):\n",
    "            train(model, optimizer, data)\n",
    "            eval_info = evaluate(model, data)\n",
    "            eval_info['epoch'] = epoch\n",
    "\n",
    "            if eval_info['val_loss'] < best_val_loss:\n",
    "                best_val_loss = eval_info['val_loss']\n",
    "                test_acc = eval_info['test_acc']\n",
    "\n",
    "            val_loss_history.append(eval_info['val_loss'])\n",
    "            if early_stopping > 0 and epoch > epochs // 2:\n",
    "                tmp = tensor(val_loss_history[-(early_stopping + 1):-1])\n",
    "                if eval_info['val_loss'] > tmp.mean().item():\n",
    "                    break\n",
    "\n",
    "        if torch.cuda.is_available():\n",
    "            torch.cuda.synchronize()\n",
    "\n",
    "        t_end = time.perf_counter()\n",
    "\n",
    "        val_losses.append(best_val_loss)\n",
    "        accs.append(test_acc)\n",
    "        durations.append(t_end - t_start)\n",
    "\n",
    "    loss, acc, duration = tensor(val_losses), tensor(accs), tensor(durations)\n",
    "\n",
    "    print('Val Loss: {:.4f}, Test Accuracy: {:.3f} ± {:.3f}, Duration: {:.3f}'.\n",
    "          format(loss.mean().item(),\n",
    "                 acc.mean().item(),\n",
    "                 acc.std().item(),\n",
    "                 duration.mean().item()))\n",
    "\n",
    "\n",
    "def train(model, optimizer, data):\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    out = model(data)\n",
    "    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "\n",
    "def evaluate(model, data):\n",
    "    model.eval()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        logits = model(data)\n",
    "\n",
    "    outs = {}\n",
    "    for key in ['train', 'val', 'test']:\n",
    "        mask = data['{}_mask'.format(key)]\n",
    "        loss = F.nll_loss(logits[mask], data.y[mask]).item()\n",
    "        pred = logits[mask].max(1)[1]\n",
    "        acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item()\n",
    "\n",
    "        outs['{}_loss'.format(key)] = loss\n",
    "        outs['{}_acc'.format(key)] = acc\n",
    "\n",
    "    return outs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Build the model with GCN on vertex classification (lecture)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Val Loss: 0.7271, Test Accuracy: 0.801 ± 0.013, Duration: 1.066\n"
     ]
    }
   ],
   "source": [
    "runs = 10\n",
    "epochs = 200\n",
    "lr = 0.01\n",
    "weight_decay = 0.0005\n",
    "early_stopping = 10\n",
    "hidden = 16\n",
    "dropout = 0.5\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "\n",
    "class Net(torch.nn.Module):\n",
    "    def __init__(self, dataset):\n",
    "        super(Net, self).__init__()\n",
    "        \n",
    "        ########################################################################\n",
    "        #      START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)             #\n",
    "        ########################################################################\n",
    "\n",
    "        self.conv1 = GCNConv(dataset.num_features, hidden)\n",
    "        self.conv2 = GCNConv(hidden, dataset.num_classes)\n",
    "        \n",
    "        ########################################################################\n",
    "        #                             END OF YOUR CODE                         #\n",
    "        ########################################################################                    \n",
    "\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        self.conv1.reset_parameters()\n",
    "        self.conv2.reset_parameters()\n",
    "\n",
    "    def forward(self, data):\n",
    "        \n",
    "        ########################################################################\n",
    "        #      START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)             #\n",
    "        ########################################################################\n",
    "\n",
    "        x, edge_index = data.x, data.edge_index\n",
    "        x = F.relu(self.conv1(x, edge_index))\n",
    "        x = F.dropout(x, p=dropout, training=self.training)\n",
    "        x = self.conv2(x, edge_index)\n",
    "        return F.log_softmax(x, dim=1)        \n",
    "        \n",
    "        ########################################################################\n",
    "        #                             END OF YOUR CODE                         #\n",
    "        ########################################################################                    \n",
    "\n",
    "    \n",
    "run(dataset, Net(dataset), runs, epochs, lr, weight_decay,\n",
    "    early_stopping)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Build models with GraphSAGE on vertex classification (lab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GraphSAGE-mean\n",
      "Val Loss: 1.3672, Test Accuracy: 0.767 ± 0.011, Duration: 0.891\n",
      "GraphSAGE-add\n",
      "Val Loss: 1.3505, Test Accuracy: 0.775 ± 0.007, Duration: 0.768\n",
      "GraphSAGE-max\n",
      "Val Loss: 1.3849, Test Accuracy: 0.751 ± 0.019, Duration: 1.198\n"
     ]
    }
   ],
   "source": [
    "## define your own model\n",
    "\n",
    "class SAGENet(torch.nn.Module):\n",
    "    def __init__(self, dataset, aggr='mean'):\n",
    "        super(SAGENet, self).__init__()\n",
    "        \n",
    "        ########################################################################\n",
    "        #      START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)             #\n",
    "        ########################################################################\n",
    "\n",
    "        self.conv1 = SAGEConv(dataset.num_features, hidden, aggr=aggr)\n",
    "        self.conv2 = SAGEConv(hidden, dataset.num_classes, aggr=aggr)        \n",
    "        \n",
    "        ########################################################################\n",
    "        #                             END OF YOUR CODE                         #\n",
    "        ########################################################################                    \n",
    "        \n",
    "        self.reset_parameters()\n",
    "        \n",
    "\n",
    "    def reset_parameters(self):\n",
    "        self.conv1.reset_parameters()\n",
    "        self.conv2.reset_parameters()\n",
    "\n",
    "    def forward(self, data):\n",
    "        \n",
    "        ########################################################################\n",
    "        #      START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)             #\n",
    "        ########################################################################\n",
    "\n",
    "        x, edge_index = data.x, data.edge_index\n",
    "        x = F.relu(self.conv1(x, edge_index))\n",
    "        x = F.dropout(x, p=dropout, training=self.training)\n",
    "        x = self.conv2(x, edge_index)\n",
    "        return F.log_softmax(x, dim=1)\n",
    "        \n",
    "        ########################################################################\n",
    "        #                             END OF YOUR CODE                         #\n",
    "        ########################################################################                    \n",
    "        \n",
    "\n",
    "    \n",
    "aggrs = ['mean', 'add', 'max']    \n",
    "\n",
    "for aggr in aggrs:\n",
    "    print('GraphSAGE-{}'.format(aggr))\n",
    "    run(dataset, SAGENet(dataset, aggr), runs, epochs, lr, weight_decay,\n",
    "        early_stopping)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph Classification\n",
    "\n",
    "While Graph Convolutional Networks (GCN) and GraphSAGE show extraordinary performance on transductive learning and inductive learning problems resepectively, then cannot learn to distinguish certain simple graph structures. **Graph Isomorphism Network (GIN)** is the state-of-the-art graph neural networks, which is provably the most expressive among the class of GNNs and is as powerful as the Weisfeiler-Lehman graph isomorphism test. \n",
    "\n",
    "(Theorem 3, How powerful are graph neural networks?) *Let $\\mathcal{A} : \\mathcal{G} → \\mathbb{R}^d$ be a GNN. With a sufficient number of GNN layers, A maps any graphs $G_1$ and $G_2$ that the Weisfeiler-Lehman test of isomorphism decides as non-isomorphic, to different embeddings if the following conditions hold:*\n",
    "\n",
    "- a) *$\\mathcal{A}$ aggregates and updates node features iteratively with*\n",
    "$$\n",
    "    h{_v}^{(k)} = \\phi \\left( h_v^{(k-1)}, f(\\{ h_u^{(k-1)}: u\\in\\mathcal{N}(v) \\}) \\right)\n",
    "$$\n",
    "*where the functions $f$, which operates on multisets, and $\\phi$ are injective.*\n",
    "- b) *$\\mathcal{A}$'s graph-level readout, which operates on the multiset ofnode features $\\{ h_v^{(k)} \\}$ , is injective.*\n",
    "\n",
    "**Graph Isomorphism Network (GIN)**, that provably satisfies the conditions in Theorem 3, is defined as:\n",
    "$$\n",
    "\\mathbf{x}^{\\prime}_i = h_{\\mathbf{\\Theta}} \\left( (1 + \\epsilon) \\cdot\n",
    "        \\mathbf{x}_i + \\mathrm{AGGR}_{j \\in \\mathcal{N}(i)} \\mathbf{x}_j \\right)\n",
    "$$\n",
    "$h_{\\mathbf{\\Theta}}$ denotes a neural network, *.i.e.* a MLP, and AGGR explicitly denotes SUM because of **higher expressive power (than MEAN/MAX)**.\n",
    "\n",
    "You should implement **GIN-0/GIn-$\\epsilon$** with **SUM/MEAN/MAX Aggregation functions** and use **MEAN Readout function** in the end of the network to obtain **graph-level representations**.\n",
    "\n",
    "<p align=\"center\"><img width=\"80%\" src=\"figures/aggr.png\"></p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.datasets import TUDataset\n",
    "from torch_geometric.utils import degree\n",
    "import torch_geometric.transforms as T\n",
    "\n",
    "\n",
    "class NormalizedDegree(object):\n",
    "    def __init__(self, mean, std):\n",
    "        self.mean = mean\n",
    "        self.std = std\n",
    "\n",
    "    def __call__(self, data):\n",
    "        deg = degree(data.edge_index[0], dtype=torch.float)\n",
    "        deg = (deg - self.mean) / self.std\n",
    "        data.x = deg.view(-1, 1)\n",
    "        return data\n",
    "\n",
    "\n",
    "def get_dataset(name, cleaned=False):\n",
    "    path = osp.join(os.getcwd(), 'data', name)\n",
    "    dataset = TUDataset(path, name, cleaned=cleaned)\n",
    "    dataset.data.edge_attr = None\n",
    "\n",
    "    if dataset.data.x is None:\n",
    "        max_degree = 0\n",
    "        degs = []\n",
    "        for data in dataset:\n",
    "            degs += [degree(data.edge_index[0], dtype=torch.long)]\n",
    "            max_degree = max(max_degree, degs[-1].max().item())\n",
    "\n",
    "        if max_degree < 1000:\n",
    "            dataset.transform = T.OneHotDegree(max_degree)\n",
    "        else:\n",
    "            deg = torch.cat(degs, dim=0).to(torch.float)\n",
    "            mean, std = deg.mean().item(), deg.std().item()\n",
    "            dataset.transform = NormalizedDegree(mean, std)\n",
    "\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name IMDB-BINARY(1000)\n",
      "Graphs 1000\n",
      "Nodes 19.773\n",
      "Edges 96.531\n",
      "Features 136\n",
      "Classes 2\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def print_dataset(dataset):\n",
    "    num_nodes = num_edges = 0\n",
    "    for data in dataset:\n",
    "        num_nodes += data.num_nodes\n",
    "        num_edges += data.num_edges\n",
    "\n",
    "    print('Name', dataset)\n",
    "    print('Graphs', len(dataset))\n",
    "    print('Nodes', num_nodes / len(dataset))\n",
    "    print('Edges', (num_edges // 2) / len(dataset))\n",
    "    print('Features', dataset.num_features)\n",
    "    print('Classes', dataset.num_classes)\n",
    "    print()\n",
    "\n",
    "\n",
    "for name in ['IMDB-BINARY']:\n",
    "    print_dataset(get_dataset(name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch import tensor\n",
    "from torch.optim import Adam\n",
    "from sklearn.model_selection import StratifiedKFold\n",
    "from torch_geometric.data import DataLoader, DenseDataLoader as DenseLoader\n",
    "\n",
    "\n",
    "def cross_validation_with_val_set(dataset, model, folds, epochs, batch_size,\n",
    "                                  lr, lr_decay_factor, lr_decay_step_size,\n",
    "                                  weight_decay, logger=None):\n",
    "\n",
    "    val_losses, accs, durations = [], [], []\n",
    "    for fold, (train_idx, test_idx,\n",
    "               val_idx) in enumerate(zip(*k_fold(dataset, folds))):\n",
    "\n",
    "        train_dataset = dataset[train_idx]\n",
    "        test_dataset = dataset[test_idx]\n",
    "        val_dataset = dataset[val_idx]\n",
    "        \n",
    "        train_loader = DataLoader(train_dataset, batch_size, shuffle=True)\n",
    "        val_loader = DataLoader(val_dataset, batch_size, shuffle=False)\n",
    "        test_loader = DataLoader(test_dataset, batch_size, shuffle=False)        \n",
    "\n",
    "        model.to(device).reset_parameters()\n",
    "        optimizer = Adam(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
    "\n",
    "        if torch.cuda.is_available():\n",
    "            torch.cuda.synchronize()\n",
    "\n",
    "        t_start = time.perf_counter()\n",
    "\n",
    "        for epoch in range(1, epochs + 1):\n",
    "            train_loss = train(model, optimizer, train_loader)\n",
    "            val_losses.append(eval_loss(model, val_loader))\n",
    "            accs.append(eval_acc(model, test_loader))\n",
    "            eval_info = {\n",
    "                'fold': fold,\n",
    "                'epoch': epoch,\n",
    "                'train_loss': train_loss,\n",
    "                'val_loss': val_losses[-1],\n",
    "                'test_acc': accs[-1],\n",
    "            }\n",
    "\n",
    "            if logger is not None:\n",
    "                logger(eval_info)\n",
    "\n",
    "            if epoch % lr_decay_step_size == 0:\n",
    "                for param_group in optimizer.param_groups:\n",
    "                    param_group['lr'] = lr_decay_factor * param_group['lr']\n",
    "\n",
    "        if torch.cuda.is_available():\n",
    "            torch.cuda.synchronize()\n",
    "\n",
    "        t_end = time.perf_counter()\n",
    "        durations.append(t_end - t_start)\n",
    "\n",
    "    loss, acc, duration = tensor(val_losses), tensor(accs), tensor(durations)\n",
    "    loss, acc = loss.view(folds, epochs), acc.view(folds, epochs)\n",
    "    loss, argmin = loss.min(dim=1)\n",
    "    acc = acc[torch.arange(folds, dtype=torch.long), argmin]\n",
    "\n",
    "    loss_mean = loss.mean().item()\n",
    "    acc_mean = acc.mean().item()\n",
    "    acc_std = acc.std().item()\n",
    "    duration_mean = duration.mean().item()\n",
    "    print('Val Loss: {:.4f}, Test Accuracy: {:.3f} ± {:.3f}, Duration: {:.3f}'.\n",
    "          format(loss_mean, acc_mean, acc_std, duration_mean))\n",
    "\n",
    "    return loss_mean, acc_mean, acc_std\n",
    "\n",
    "\n",
    "def k_fold(dataset, folds):\n",
    "    skf = StratifiedKFold(folds, shuffle=True, random_state=12345)\n",
    "\n",
    "    test_indices, train_indices = [], []\n",
    "    for _, idx in skf.split(torch.zeros(len(dataset)), dataset.data.y):\n",
    "        test_indices.append(torch.from_numpy(idx))\n",
    "\n",
    "    val_indices = [test_indices[i - 1] for i in range(folds)]\n",
    "\n",
    "    for i in range(folds):\n",
    "        train_mask = torch.ones(len(dataset), dtype=torch.bool)\n",
    "        train_mask[test_indices[i]] = 0\n",
    "        train_mask[val_indices[i]] = 0\n",
    "        train_indices.append(train_mask.nonzero().view(-1))\n",
    "\n",
    "    return train_indices, test_indices, val_indices\n",
    "\n",
    "\n",
    "def num_graphs(data):\n",
    "    if data.batch is not None:\n",
    "        return data.num_graphs\n",
    "    else:\n",
    "        return data.x.size(0)\n",
    "\n",
    "\n",
    "def train(model, optimizer, loader):\n",
    "    model.train()\n",
    "\n",
    "    total_loss = 0\n",
    "    for data in loader:\n",
    "        optimizer.zero_grad()\n",
    "        data = data.to(device)\n",
    "        out = model(data)\n",
    "        loss = F.nll_loss(out, data.y.view(-1))\n",
    "        loss.backward()\n",
    "        total_loss += loss.item() * num_graphs(data)\n",
    "        optimizer.step()\n",
    "    return total_loss / len(loader.dataset)\n",
    "\n",
    "\n",
    "def eval_acc(model, loader):\n",
    "    model.eval()\n",
    "\n",
    "    correct = 0\n",
    "    for data in loader:\n",
    "        data = data.to(device)\n",
    "        with torch.no_grad():\n",
    "            pred = model(data).max(1)[1]\n",
    "        correct += pred.eq(data.y.view(-1)).sum().item()\n",
    "    return correct / len(loader.dataset)\n",
    "\n",
    "\n",
    "def eval_loss(model, loader):\n",
    "    model.eval()\n",
    "\n",
    "    loss = 0\n",
    "    for data in loader:\n",
    "        data = data.to(device)\n",
    "        with torch.no_grad():\n",
    "            out = model(data)\n",
    "        loss += F.nll_loss(out, data.y.view(-1), reduction='sum').item()\n",
    "    return loss / len(loader.dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Implement Graph Isomorphism Network (lab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.nn import Linear, Sequential, ReLU, BatchNorm1d as BN\n",
    "from torch_geometric.nn import global_mean_pool, MessagePassing\n",
    "from torch_geometric.utils import remove_self_loops\n",
    "\n",
    "def reset(nn):\n",
    "    def _reset(item):\n",
    "        if hasattr(item, 'reset_parameters'):\n",
    "            item.reset_parameters()\n",
    "\n",
    "    if nn is not None:\n",
    "        if hasattr(nn, 'children') and len(list(nn.children())) > 0:\n",
    "            for item in nn.children():\n",
    "                _reset(item)\n",
    "        else:\n",
    "            _reset(nn)\n",
    "\n",
    "\n",
    "class GINConv(MessagePassing):\n",
    "    def __init__(self, nn, eps=0, train_eps=False, **kwargs):\n",
    "        super(GINConv, self).__init__(aggr='add', **kwargs)\n",
    "        self.nn = nn\n",
    "        self.initial_eps = eps\n",
    "        if train_eps:\n",
    "            self.eps = torch.nn.Parameter(torch.Tensor([eps]))\n",
    "        else:\n",
    "            self.register_buffer('eps', torch.Tensor([eps]))\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        reset(self.nn)\n",
    "        self.eps.data.fill_(self.initial_eps)\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        \"\"\"\"\"\"\n",
    "        \n",
    "        ########################################################################\n",
    "        #      START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)             #\n",
    "        ########################################################################\n",
    "\n",
    "        edge_index, _ = remove_self_loops(edge_index)\n",
    "        out = self.nn((1 + self.eps) * x + self.propagate(edge_index, x=x))\n",
    "        return out        \n",
    "        \n",
    "        ########################################################################\n",
    "        #                             END OF YOUR CODE                         #\n",
    "        ########################################################################            \n",
    "        \n",
    "\n",
    "\n",
    "    def message(self, x_j):\n",
    "        \n",
    "        ########################################################################\n",
    "        #      START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)             #\n",
    "        ########################################################################\n",
    "\n",
    "        return x_j        \n",
    "        \n",
    "        ########################################################################\n",
    "        #                             END OF YOUR CODE                         #\n",
    "        ########################################################################                    \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Implement models with GIN-0/GIN-$\\epsilon$ to perform graph classificaiton on IMDB-binary dataset (lab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GIN0(torch.nn.Module):\n",
    "    def __init__(self, dataset, num_layers, hidden):\n",
    "        super(GIN0, self).__init__()\n",
    "        self.conv1 = GINConv(Sequential(\n",
    "            Linear(dataset.num_features, hidden),\n",
    "            ReLU(),\n",
    "            Linear(hidden, hidden),\n",
    "            ReLU(),\n",
    "            BN(hidden),\n",
    "        ),\n",
    "                             train_eps=False)\n",
    "        self.convs = torch.nn.ModuleList()\n",
    "        for i in range(num_layers - 1):\n",
    "            self.convs.append(\n",
    "                GINConv(Sequential(\n",
    "                    Linear(hidden, hidden),\n",
    "                    ReLU(),\n",
    "                    Linear(hidden, hidden),\n",
    "                    ReLU(),\n",
    "                    BN(hidden),\n",
    "                ),\n",
    "                        train_eps=False))\n",
    "        self.lin1 = Linear(hidden, hidden)\n",
    "        self.lin2 = Linear(hidden, dataset.num_classes)\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        self.conv1.reset_parameters()\n",
    "        for conv in self.convs:\n",
    "            conv.reset_parameters()\n",
    "        self.lin1.reset_parameters()\n",
    "        self.lin2.reset_parameters()\n",
    "\n",
    "    def forward(self, data):\n",
    "        \n",
    "        ########################################################################\n",
    "        #      START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)             #\n",
    "        ########################################################################\n",
    "\n",
    "        x, edge_index, batch = data.x, data.edge_index, data.batch\n",
    "        x = self.conv1(x, edge_index)\n",
    "        for conv in self.convs:\n",
    "            x = conv(x, edge_index)\n",
    "        x = global_mean_pool(x, batch)\n",
    "        x = F.relu(self.lin1(x))\n",
    "        x = F.dropout(x, p=0.5, training=self.training)\n",
    "        x = self.lin2(x)\n",
    "        return F.log_softmax(x, dim=-1)\n",
    "        \n",
    "        ########################################################################\n",
    "        #                             END OF YOUR CODE                         #\n",
    "        ########################################################################                    \n",
    "        \n",
    "\n",
    "\n",
    "\n",
    "class GIN(torch.nn.Module):\n",
    "    def __init__(self, dataset, num_layers, hidden):\n",
    "        super(GIN, self).__init__()\n",
    "        self.conv1 = GINConv(Sequential(\n",
    "            Linear(dataset.num_features, hidden),\n",
    "            ReLU(),\n",
    "            Linear(hidden, hidden),\n",
    "            ReLU(),\n",
    "            BN(hidden),\n",
    "        ),\n",
    "                             train_eps=True)\n",
    "        self.convs = torch.nn.ModuleList()\n",
    "        for i in range(num_layers - 1):\n",
    "            self.convs.append(\n",
    "                GINConv(Sequential(\n",
    "                    Linear(hidden, hidden),\n",
    "                    ReLU(),\n",
    "                    Linear(hidden, hidden),\n",
    "                    ReLU(),\n",
    "                    BN(hidden),\n",
    "                ),\n",
    "                        train_eps=True))\n",
    "        self.lin1 = Linear(hidden, hidden)\n",
    "        self.lin2 = Linear(hidden, dataset.num_classes)\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        self.conv1.reset_parameters()\n",
    "        for conv in self.convs:\n",
    "            conv.reset_parameters()\n",
    "        self.lin1.reset_parameters()\n",
    "        self.lin2.reset_parameters()\n",
    "\n",
    "    def forward(self, data):\n",
    "        \n",
    "        ########################################################################\n",
    "        #      START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)             #\n",
    "        ########################################################################\n",
    "\n",
    "        x, edge_index, batch = data.x, data.edge_index, data.batch\n",
    "        x = self.conv1(x, edge_index)\n",
    "        for conv in self.convs:\n",
    "            x = conv(x, edge_index)\n",
    "        x = global_mean_pool(x, batch)\n",
    "        x = F.relu(self.lin1(x))\n",
    "        x = F.dropout(x, p=0.5, training=self.training)\n",
    "        x = self.lin2(x)\n",
    "        return F.log_softmax(x, dim=-1)\n",
    "        \n",
    "        \n",
    "        ########################################################################\n",
    "        #                             END OF YOUR CODE                         #\n",
    "        ########################################################################                    \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-----\n",
      "IMDB-BINARY - GIN0\n",
      "Val Loss: 0.4679, Test Accuracy: 0.720 ± 0.049, Duration: 37.890\n",
      "Best result - 0.720 ± 0.049\n",
      "-----\n",
      "IMDB-BINARY - GIN\n"
     ]
    }
   ],
   "source": [
    "from itertools import product\n",
    "\n",
    "epochs = 100\n",
    "batch_size = 128\n",
    "lr = 0.01\n",
    "lr_decay_factor = 0.5\n",
    "lr_decay_step_size = 50\n",
    "\n",
    "layers = [5]\n",
    "hiddens = [64]\n",
    "datasets = ['IMDB-BINARY']\n",
    "nets = [\n",
    "    GIN0,\n",
    "    GIN,\n",
    "]\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "def logger(info):\n",
    "    fold, epoch = info['fold'] + 1, info['epoch']\n",
    "    val_loss, test_acc = info['val_loss'], info['test_acc']\n",
    "    print('{:02d}/{:03d}: Val Loss: {:.4f}, Test Accuracy: {:.3f}'.format(\n",
    "        fold, epoch, val_loss, test_acc))\n",
    "\n",
    "\n",
    "results = []\n",
    "for dataset_name, Net in product(datasets, nets):\n",
    "    best_result = (float('inf'), 0, 0)  # (loss, acc, std)\n",
    "    print('-----\\n{} - {}'.format(dataset_name, Net.__name__))\n",
    "    for num_layers, hidden in product(layers, hiddens):\n",
    "        dataset = get_dataset(dataset_name)\n",
    "        model = Net(dataset, num_layers, hidden)\n",
    "        loss, acc, std = cross_validation_with_val_set(\n",
    "            dataset,\n",
    "            model,\n",
    "            folds=10,\n",
    "            epochs=epochs,\n",
    "            batch_size=batch_size,\n",
    "            lr=lr,\n",
    "            lr_decay_factor=lr_decay_factor,\n",
    "            lr_decay_step_size=lr_decay_step_size,\n",
    "            weight_decay=0,\n",
    "            logger=None,\n",
    "        )\n",
    "        if loss < best_result[0]:\n",
    "            best_result = (loss, acc, std)\n",
    "\n",
    "    desc = '{:.3f} ± {:.3f}'.format(best_result[1], best_result[2])\n",
    "    print('Best result - {}'.format(desc))\n",
    "    results += ['{} - {}: {}'.format(dataset_name, model, desc)]\n",
    "print('-----\\n{}'.format('\\n'.join(results)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-----\n",
      "IMDB-BINARY - GIN0\n",
      "Val Loss: 0.4650, Test Accuracy: 0.722 ± 0.048, Duration: 29.156\n",
      "Best result - 0.722 ± 0.048\n",
      "-----\n",
      "IMDB-BINARY - GIN\n",
      "Val Loss: 0.4712, Test Accuracy: 0.707 ± 0.053, Duration: 30.204\n",
      "Best result - 0.707 ± 0.053\n",
      "-----\n",
      "IMDB-BINARY - GIN0(\n",
      "  (conv1): GINConv(nn=Sequential(\n",
      "    (0): Linear(in_features=136, out_features=64, bias=True)\n",
      "    (1): ReLU()\n",
      "    (2): Linear(in_features=64, out_features=64, bias=True)\n",
      "    (3): ReLU()\n",
      "    (4): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  ))\n",
      "  (convs): ModuleList(\n",
      "    (0): GINConv(nn=Sequential(\n",
      "      (0): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (1): ReLU()\n",
      "      (2): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (3): ReLU()\n",
      "      (4): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    ))\n",
      "    (1): GINConv(nn=Sequential(\n",
      "      (0): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (1): ReLU()\n",
      "      (2): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (3): ReLU()\n",
      "      (4): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    ))\n",
      "    (2): GINConv(nn=Sequential(\n",
      "      (0): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (1): ReLU()\n",
      "      (2): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (3): ReLU()\n",
      "      (4): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    ))\n",
      "    (3): GINConv(nn=Sequential(\n",
      "      (0): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (1): ReLU()\n",
      "      (2): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (3): ReLU()\n",
      "      (4): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    ))\n",
      "  )\n",
      "  (lin1): Linear(in_features=64, out_features=64, bias=True)\n",
      "  (lin2): Linear(in_features=64, out_features=2, bias=True)\n",
      "): 0.722 ± 0.048\n",
      "IMDB-BINARY - GIN(\n",
      "  (conv1): GINConv(nn=Sequential(\n",
      "    (0): Linear(in_features=136, out_features=64, bias=True)\n",
      "    (1): ReLU()\n",
      "    (2): Linear(in_features=64, out_features=64, bias=True)\n",
      "    (3): ReLU()\n",
      "    (4): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  ))\n",
      "  (convs): ModuleList(\n",
      "    (0): GINConv(nn=Sequential(\n",
      "      (0): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (1): ReLU()\n",
      "      (2): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (3): ReLU()\n",
      "      (4): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    ))\n",
      "    (1): GINConv(nn=Sequential(\n",
      "      (0): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (1): ReLU()\n",
      "      (2): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (3): ReLU()\n",
      "      (4): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    ))\n",
      "    (2): GINConv(nn=Sequential(\n",
      "      (0): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (1): ReLU()\n",
      "      (2): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (3): ReLU()\n",
      "      (4): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    ))\n",
      "    (3): GINConv(nn=Sequential(\n",
      "      (0): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (1): ReLU()\n",
      "      (2): Linear(in_features=64, out_features=64, bias=True)\n",
      "      (3): ReLU()\n",
      "      (4): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    ))\n",
      "  )\n",
      "  (lin1): Linear(in_features=64, out_features=64, bias=True)\n",
      "  (lin2): Linear(in_features=64, out_features=2, bias=True)\n",
      "): 0.707 ± 0.053\n"
     ]
    }
   ],
   "source": [
    "from itertools import product\n",
    "\n",
    "epochs = 100\n",
    "batch_size = 128\n",
    "lr = 0.01\n",
    "lr_decay_factor = 0.5\n",
    "lr_decay_step_size = 50\n",
    "\n",
    "layers = [5]\n",
    "hiddens = [64]\n",
    "datasets = ['IMDB-BINARY']\n",
    "nets = [\n",
    "    GIN0,\n",
    "    GIN,\n",
    "]\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "def logger(info):\n",
    "    fold, epoch = info['fold'] + 1, info['epoch']\n",
    "    val_loss, test_acc = info['val_loss'], info['test_acc']\n",
    "    print('{:02d}/{:03d}: Val Loss: {:.4f}, Test Accuracy: {:.3f}'.format(\n",
    "        fold, epoch, val_loss, test_acc))\n",
    "\n",
    "\n",
    "results = []\n",
    "for dataset_name, Net in product(datasets, nets):\n",
    "    best_result = (float('inf'), 0, 0)  # (loss, acc, std)\n",
    "    print('-----\\n{} - {}'.format(dataset_name, Net.__name__))\n",
    "    for num_layers, hidden in product(layers, hiddens):\n",
    "        dataset = get_dataset(dataset_name)\n",
    "        model = Net(dataset, num_layers, hidden)\n",
    "        loss, acc, std = cross_validation_with_val_set(\n",
    "            dataset,\n",
    "            model,\n",
    "            folds=10,\n",
    "            epochs=epochs,\n",
    "            batch_size=batch_size,\n",
    "            lr=lr,\n",
    "            lr_decay_factor=lr_decay_factor,\n",
    "            lr_decay_step_size=lr_decay_step_size,\n",
    "            weight_decay=0,\n",
    "            logger=None,\n",
    "        )\n",
    "        if loss < best_result[0]:\n",
    "            best_result = (loss, acc, std)\n",
    "\n",
    "    desc = '{:.3f} ± {:.3f}'.format(best_result[1], best_result[2])\n",
    "    print('Best result - {}'.format(desc))\n",
    "    results += ['{} - {}: {}'.format(dataset_name, model, desc)]\n",
    "print('-----\\n{}'.format('\\n'.join(results)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pyg",
   "language": "python",
   "name": "pyg"
  },
  "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}