ModernGAP.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Lua-Nova/Modern-GAP-GNN/blob/main/ModernGAP.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "aFQCNsDWU5mf",
        "outputId": "1964c151-8fb0-4a5c-c781-79f5ca3ba9c8"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Looking in links: https://data.pyg.org/whl/torch-1.12.0+cpu.html\n",
            "Collecting torch-scatter\n",
            "  Downloading https://data.pyg.org/whl/torch-1.12.0%2Bcpu/torch_scatter-2.0.9-cp37-cp37m-linux_x86_64.whl (286 kB)\n",
            "\u001b[K     |████████████████████████████████| 286 kB 4.1 MB/s \n",
            "\u001b[?25hCollecting torch-sparse\n",
            "  Downloading https://data.pyg.org/whl/torch-1.12.0%2Bcpu/torch_sparse-0.6.15-cp37-cp37m-linux_x86_64.whl (641 kB)\n",
            "\u001b[K     |████████████████████████████████| 641 kB 9.1 MB/s \n",
            "\u001b[?25hCollecting torch-cluster\n",
            "  Downloading https://data.pyg.org/whl/torch-1.12.0%2Bcpu/torch_cluster-1.6.0-cp37-cp37m-linux_x86_64.whl (311 kB)\n",
            "\u001b[K     |████████████████████████████████| 311 kB 45.3 MB/s \n",
            "\u001b[?25hCollecting torch-spline-conv\n",
            "  Downloading https://data.pyg.org/whl/torch-1.12.0%2Bcpu/torch_spline_conv-1.2.1-cp37-cp37m-linux_x86_64.whl (121 kB)\n",
            "\u001b[K     |████████████████████████████████| 121 kB 8.5 MB/s \n",
            "\u001b[?25hCollecting torch-geometric\n",
            "  Downloading torch_geometric-2.1.0.post1.tar.gz (467 kB)\n",
            "\u001b[K     |████████████████████████████████| 467 kB 4.5 MB/s \n",
            "\u001b[?25hRequirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from torch-sparse) (1.7.3)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (4.64.1)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (1.21.6)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (2.11.3)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (2.23.0)\n",
            "Requirement already satisfied: pyparsing in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (3.0.9)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (1.0.2)\n",
            "Requirement already satisfied: MarkupSafe>=0.23 in /usr/local/lib/python3.7/dist-packages (from jinja2->torch-geometric) (2.0.1)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->torch-geometric) (3.0.4)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->torch-geometric) (1.24.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->torch-geometric) (2.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->torch-geometric) (2022.9.24)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from scikit-learn->torch-geometric) (3.1.0)\n",
            "Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.7/dist-packages (from scikit-learn->torch-geometric) (1.2.0)\n",
            "Building wheels for collected packages: torch-geometric\n",
            "  Building wheel for torch-geometric (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for torch-geometric: filename=torch_geometric-2.1.0.post1-py3-none-any.whl size=689859 sha256=02b187fc7f1a045e1e6c089b1bc5606b823f1db0ae2f414d08901f02d6f6a927\n",
            "  Stored in directory: /root/.cache/pip/wheels/d1/cb/43/f7f2e472de4d7cff31bceddadc36d634e1e545fbc17961c282\n",
            "Successfully built torch-geometric\n",
            "Installing collected packages: torch-spline-conv, torch-sparse, torch-scatter, torch-geometric, torch-cluster\n",
            "Successfully installed torch-cluster-1.6.0 torch-geometric-2.1.0.post1 torch-scatter-2.0.9 torch-sparse-0.6.15 torch-spline-conv-1.2.1\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting opacus\n",
            "  Downloading opacus-1.3.0-py3-none-any.whl (216 kB)\n",
            "\u001b[K     |████████████████████████████████| 216 kB 4.6 MB/s \n",
            "\u001b[?25hCollecting functorch\n",
            "  Downloading functorch-1.13.0-py2.py3-none-any.whl (2.1 kB)\n",
            "Requirement already satisfied: opt-einsum>=3.3.0 in /usr/local/lib/python3.7/dist-packages (from opacus) (3.3.0)\n",
            "Requirement already satisfied: torch>=1.8 in /usr/local/lib/python3.7/dist-packages (from opacus) (1.12.1+cu113)\n",
            "Requirement already satisfied: numpy>=1.15 in /usr/local/lib/python3.7/dist-packages (from opacus) (1.21.6)\n",
            "Requirement already satisfied: scipy>=1.2 in /usr/local/lib/python3.7/dist-packages (from opacus) (1.7.3)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch>=1.8->opacus) (4.1.1)\n",
            "Collecting torch>=1.8\n",
            "  Downloading torch-1.13.0-cp37-cp37m-manylinux1_x86_64.whl (890.2 MB)\n",
            "\u001b[K     |██████████████████████████████  | 834.1 MB 1.2 MB/s eta 0:00:49tcmalloc: large alloc 1147494400 bytes == 0x3b50000 @  0x7f9397aa2615 0x58ead6 0x4f355e 0x4d222f 0x51041f 0x5b4ee6 0x58ff2e 0x510325 0x5b4ee6 0x58ff2e 0x50d482 0x4d00fb 0x50cb8d 0x4d00fb 0x50cb8d 0x4d00fb 0x50cb8d 0x4bac0a 0x538a76 0x590ae5 0x510280 0x5b4ee6 0x58ff2e 0x50d482 0x5b4ee6 0x58ff2e 0x50c4fc 0x58fd37 0x50ca37 0x5b4ee6 0x58ff2e\n",
            "\u001b[K     |████████████████████████████████| 890.2 MB 6.1 kB/s \n",
            "\u001b[?25hCollecting nvidia-cublas-cu11==11.10.3.66\n",
            "  Downloading nvidia_cublas_cu11-11.10.3.66-py3-none-manylinux1_x86_64.whl (317.1 MB)\n",
            "\u001b[K     |████████████████████████████████| 317.1 MB 30 kB/s \n",
            "\u001b[?25hCollecting nvidia-cuda-nvrtc-cu11==11.7.99\n",
            "  Downloading nvidia_cuda_nvrtc_cu11-11.7.99-2-py3-none-manylinux1_x86_64.whl (21.0 MB)\n",
            "\u001b[K     |████████████████████████████████| 21.0 MB 1.2 MB/s \n",
            "\u001b[?25hCollecting nvidia-cudnn-cu11==8.5.0.96\n",
            "  Downloading nvidia_cudnn_cu11-8.5.0.96-2-py3-none-manylinux1_x86_64.whl (557.1 MB)\n",
            "\u001b[K     |████████████████████████████████| 557.1 MB 12 kB/s \n",
            "\u001b[?25hCollecting nvidia-cuda-runtime-cu11==11.7.99\n",
            "  Downloading nvidia_cuda_runtime_cu11-11.7.99-py3-none-manylinux1_x86_64.whl (849 kB)\n",
            "\u001b[K     |████████████████████████████████| 849 kB 42.7 MB/s \n",
            "\u001b[?25hRequirement already satisfied: wheel in /usr/local/lib/python3.7/dist-packages (from nvidia-cublas-cu11==11.10.3.66->torch>=1.8->opacus) (0.38.4)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from nvidia-cublas-cu11==11.10.3.66->torch>=1.8->opacus) (57.4.0)\n",
            "Installing collected packages: nvidia-cublas-cu11, nvidia-cudnn-cu11, nvidia-cuda-runtime-cu11, nvidia-cuda-nvrtc-cu11, torch, functorch, opacus\n",
            "  Attempting uninstall: torch\n",
            "    Found existing installation: torch 1.12.1+cu113\n",
            "    Uninstalling torch-1.12.1+cu113:\n",
            "      Successfully uninstalled torch-1.12.1+cu113\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "torchvision 0.13.1+cu113 requires torch==1.12.1, but you have torch 1.13.0 which is incompatible.\n",
            "torchtext 0.13.1 requires torch==1.12.1, but you have torch 1.13.0 which is incompatible.\n",
            "torchaudio 0.12.1+cu113 requires torch==1.12.1, but you have torch 1.13.0 which is incompatible.\u001b[0m\n",
            "Successfully installed functorch-1.13.0 nvidia-cublas-cu11-11.10.3.66 nvidia-cuda-nvrtc-cu11-11.7.99 nvidia-cuda-runtime-cu11-11.7.99 nvidia-cudnn-cu11-8.5.0.96 opacus-1.3.0 torch-1.13.0\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "torch"
                ]
              }
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "import torch\n",
        "if torch.cuda.is_available():\n",
        "  #NVIDIA GPU version\n",
        "  %pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f f'https://data.pyg.org/whl/torch-1.12.0+{cutorch.version.cuda.replace('.','')}.html'\n",
        "else:\n",
        "  #CPU version\n",
        "  %pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-1.12.0+cpu.html\n",
        "%pip install opacus\n",
        "\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fBxqGGgYU6k-"
      },
      "outputs": [],
      "source": [
        "import torch.nn as nn\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from torch_geometric.nn import Sequential, GCNConv\n",
        "import opacus as op"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vb1hVZklFoDP"
      },
      "source": [
        "## Encoder Module"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1ScntAVtWL3M"
      },
      "outputs": [],
      "source": [
        "torch.manual_seed(11)\n",
        "# create classes for layers that are used a lot to avoid repeating code\n",
        "\n",
        "class MLP(nn.Module):\n",
        "  # e.g. dimensions = [50,40,30,20]\n",
        "    def __init__(self, dimensions):\n",
        "        super().__init__()\n",
        "        self.flatten = nn.Flatten()\n",
        "        layers = []\n",
        "        for i in range(len(dimensions)-1):\n",
        "          layers.append(nn.Linear(dimensions[i], dimensions[i+1]))\n",
        "          layers.append(nn.SELU(inplace=True))\n",
        "\n",
        "        self.linear_relu_stack = nn.Sequential(*layers)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.flatten(x)\n",
        "        logits = self.linear_relu_stack(x)\n",
        "        return logits"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GF264FWv0pJI"
      },
      "source": [
        "## PMA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OMzLWqwF03Uu"
      },
      "outputs": [],
      "source": [
        "class PMA(nn.Module):\n",
        "    # A - adjacency matrix     TODO: this should not be given to the module itself, it should access it in training (or from the graph dataset)\n",
        "    # num_hops - the number of hops covered by this GNN\n",
        "    def __init__(self, num_hops, sigma):\n",
        "        super().__init__()\n",
        "        # TODO: Figure out if you should tranpose this\n",
        "        # self.A_transpose = torch.transpose(A, 0,1)\n",
        "        self.num_hops = num_hops\n",
        "        self.sigma = sigma\n",
        "        self.A = None\n",
        "    \n",
        "    def set_adj_mat(self, A):\n",
        "        self.A = A\n",
        "\n",
        "    def forward(self, x):\n",
        "        if self.A is None:\n",
        "          raise RuntimeError(\"Set adjacency matrix before running\")\n",
        "        out = [torch.nn.functional.normalize(x, dim=1)]\n",
        "        for k in range(self.num_hops):\n",
        "            aggr = torch.mm(self.A, out[-1])\n",
        "            noised = aggr + torch.normal(torch.zeros(aggr.size()), std=self.sigma)\n",
        "            normalized = torch.nn.functional.normalize(noised, dim=1)\n",
        "            out.append(normalized)\n",
        "        return torch.stack(out)\n",
        "        # return torch.nn.functional.normalize(x, dim=1)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class PMAT(nn.Module):\n",
        "    def __init__(self, num_hops, transform_dimensions, attention_dimensions):\n",
        "        self.num_hops = num_hops\n",
        "        self.A = None\n",
        "        self.transforms = nn.ModuleList()\n",
        "        self.attentions = nn.ModuleList()\n",
        "        for i in range(num_hops):\n",
        "          self.transforms.append(nn.Linear(*transform_dimensions)) # Only 1 layer transformation\n",
        "          self.attentions.append(MLP(attention_dimensions))\n",
        "\n",
        "    def set_adj_mat(self, A):\n",
        "        pass\n",
        "\n",
        "    def forward(self, x):\n",
        "        if self.A is None:\n",
        "          raise RuntimeError(\"Set adjacency matrix before running\")\n",
        "        out = [torch.nn.functional.normalize(x, dim=1)]\n",
        "        h = x\n",
        "        for k in range(self.num_hops):\n",
        "\n",
        "            # torch.sparse.sum(torch.index_select(A, 0, torch.tensor([0, 1, 2, 3])))\n",
        "\n",
        "            neighbour_filter = neighbour_indices[node] # self.A[node, :]\n",
        "\n",
        "            # e = torch.zeros(self.dimensions.size(dim=1))\n",
        "\n",
        "            h = self.transforms[k](h)\n",
        "\n",
        "            e = self.attentions[k](h[edge_index.T])\n",
        "\n",
        "            \n",
        "\n",
        "            # edge_indices[node]\n",
        "            # node_indices_of_neighbours[node]\n",
        "\n",
        "            # node_indices_of_neighbours[node] = edge_index[edge_indices[node]][:, 1]\n",
        "\n",
        "            alpha[edge_indices[node]] = torch.Softmax(e[edge_indices[node]], dim=0)\n",
        "\n",
        "            # alpha[edge_indices[nodes]] = torch.exp(e[edge_indices[nodes]]) / torch.sum(torch.exp(e[edge_indices[nodes]]))\n",
        "\n",
        "            h_prime[nodes] = torch.sum(alpha[edge_indices[nodes]] * h[node_indices_of_neighbours[nodes]])\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "            aggr = torch.mm(self.A, out[-1])\n",
        "            noised = aggr + torch.normal(torch.zeros(aggr.size()), std=self.sigma)\n",
        "            normalized = torch.nn.functional.normalize(noised, dim=1)\n",
        "            out.append(normalized)\n",
        "        return torch.stack(out)\n"
      ],
      "metadata": {
        "id": "GSphEJ9xipcn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CZnwBQti_gK-"
      },
      "outputs": [],
      "source": [
        "# # TEMP CODE\n",
        "# smoothing = 0.2\n",
        "# A = torch.tensor([[1.,smoothing,smoothing],\n",
        "#                   [smoothing,1.,smoothing],\n",
        "#                   [smoothing,smoothing,1.]])\n",
        "# x = torch.tensor([[1.,0.,0.],[0.,1.,0.],[0.,0.,1.]])\n",
        "# pma = PMA(10, 1)\n",
        "# tensor = pma(x, A)\n",
        "# tensor = tensor.cpu().numpy()\n",
        "\n",
        "# # plt.figure(figsize=(16,7))\n",
        "# # plt.imshow(tensor)\n",
        "# # plt.show()\n",
        "#         # [encoder, pma, element_wise_mlp, combine, mlp]\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Oaw4OM5_2Tmk",
        "outputId": "bcf3b8a0-f0f0-456f-b8b6-342b10deab0e"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[10.7376, 11.9459,  9.3005],\n",
              "        [ 8.6977,  9.4867,  9.7304]])"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ],
      "source": [
        "10*torch.ones((2, 3)) + torch.normal(torch.zeros((2, 3)), std=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ty_dao27oEAB"
      },
      "source": [
        "## Classification Module\n",
        "NOTE: \n",
        "\n",
        "MLP base: The first MLP in the cassification module. \n",
        "\n",
        "MLP head: The last MLP and takes the combined output of all MLP base."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7dLm3Q3HofT6"
      },
      "outputs": [],
      "source": [
        "class Classification(nn.Module):\n",
        "    # num_hops - the number of hops covered by this GNN\n",
        "    # encoder_dimensions - the MLP dimensions of each base MLP\n",
        "    # head_dimensions - the dimensions of the head MLP\n",
        "    def __init__(self, num_hops, encoder_dimensions, head_dimensions):\n",
        "        super().__init__()\n",
        "        self.base_mlps = nn.ModuleList()\n",
        "        for i in range(num_hops+1):\n",
        "          self.base_mlps.append(MLP(encoder_dimensions))\n",
        "        self.head_mlp = MLP(head_dimensions) # TODO: should this be softmax? I think we add a softmax for classification tasks. We can test if it works better\n",
        "    \n",
        "    def forward(self, cache):\n",
        "        # forward through bases\n",
        "        out = []\n",
        "        for i in range(len(self.base_mlps)):\n",
        "          encoding = self.base_mlps[i](cache[i,:,:])\n",
        "          out.append(encoding) # add corresponding encoding\n",
        "          # TEMP\n",
        "          # out.append(cache[i, :, :])\n",
        "        # combine (use concatenation)\n",
        "        combined_x = torch.cat(out, dim=1)\n",
        "        # forward through head\n",
        "        return self.head_mlp(combined_x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "evGOLmy-UrMs"
      },
      "outputs": [],
      "source": [
        "class GAP(nn.Module):\n",
        "  # encoder - pretrained encoder module\n",
        "  # pma - PMA module\n",
        "  # classification - classification module\n",
        "  def __init__(self, encoder, pma, classification): # TODO: decide whether we should recieve the models as parameters\n",
        "    super().__init__()\n",
        "    self.encoder = encoder\n",
        "    self.encoder.requires_grad=False\n",
        "    self.pma = pma\n",
        "    self.classification = classification\n",
        "\n",
        "  def set_adj_mat(self, A):\n",
        "    self.pma.set_adj_mat(A)\n",
        "\n",
        "  def forward(self, x):\n",
        "    # initial node encoding\n",
        "    x_encoded = self.encoder(x)\n",
        "    # aggregation module\n",
        "    cache = self.pma(x_encoded) \n",
        "    # classification\n",
        "    return self.classification(cache) \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WebCohAOIX0P"
      },
      "source": [
        "##Hyperparameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G_Pf4P_bIZ5S",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "23420594-a72b-46fe-b8c9-54b166aa90f6"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "sigma: 0\n"
          ]
        }
      ],
      "source": [
        "node_level = True\n",
        "\n",
        "# Edge level DP\n",
        "epsilon, delta, alpha = 1000, 0.1, 1\n",
        "# specify specific epsilon_1, epsilon_5 for node-level and then just do a subtraction, and calculate sigma from remaining epsilon\n",
        "K_hop = 5\n",
        "# sigma = 1 / np.max(np.roots([K_hop/2, np.sqrt(2*K_hop*np.log(1/delta)), -epsilon]))\n",
        "sigma = 0\n",
        "# Node level DP\n",
        "if (node_level):\n",
        "  pass\n",
        "  # How do we calculate this?\n",
        "data = \"reddit\"\n",
        "\n",
        "print(\"sigma:\", sigma)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E7oEXJ63oyl-"
      },
      "source": [
        "## Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IPX-rb88ukrc"
      },
      "outputs": [],
      "source": [
        "from torch_geometric.data import Data\n",
        "\n",
        "# this method partitions based on nodes (so edges between splits are not used)\n",
        "def train_test_split(dataset, test_ratio):\n",
        "    X, y, edge_index= dataset.x, dataset.y, dataset.edge_index\n",
        "    shuffle_ordering = torch.randperm(X.size(dim=0))\n",
        "\n",
        "    edge_mapping = torch.zeros(X.size(dim=0), dtype=torch.long)\n",
        "    edge_mapping[shuffle_ordering] = torch.arange(X.size(dim=0))\n",
        "\n",
        "    X = X[shuffle_ordering]\n",
        "    y = y[shuffle_ordering]\n",
        "    edge_index = edge_mapping[edge_index]\n",
        "\n",
        "    mask = torch.zeros(X.size(dim=0), dtype=torch.bool)\n",
        "    train_slice = int((1-test_ratio)*X.size(dim=0))\n",
        "    mask[:train_slice] = True\n",
        "\n",
        "    X_train = X[mask]\n",
        "    X_test = X[~mask]\n",
        "\n",
        "    y_train = y[mask]\n",
        "    y_test = y[~mask]\n",
        "\n",
        "    edge_index_train = edge_index[:, torch.logical_and(*mask[edge_index])]\n",
        "    edge_index_test = edge_index[:, torch.logical_and(*~mask[edge_index])] - train_slice\n",
        "\n",
        "    return Data(x=X_train, y=y_train, edge_index=edge_index_train), \\\n",
        "           Data(x=X_test, y=y_test, edge_index=edge_index_test)\n",
        "\n",
        "\n",
        "# returns filtered edge index, first removes edges that have removed src or dst nodes, then shifts indices of remained src/dst nodes\n",
        "def filter_edge_index(edge_index, filter):\n",
        "\n",
        "    node_indices = torch.arange(filter.size(dim=0))[filter]\n",
        "    edge_mapping = torch.zeros(filter.size(dim=0), dtype=torch.long)\n",
        "    edge_mapping[node_indices] = torch.arange(node_indices.size(dim=0))\n",
        "\n",
        "\n",
        "    edge_index = edge_index.to(torch.long)\n",
        "    edge_filter = torch.logical_and(*filter[edge_index])\n",
        "    return edge_mapping[edge_index[:, edge_filter]]\n",
        "\n",
        "def prepare_dataset(dataset, threshold):\n",
        "    X, y, edge_index = dataset.x, dataset.y, dataset.edge_index\n",
        "\n",
        "    # remove labels with less examples than threshold\n",
        "    index_map = torch.zeros(y.size())\n",
        "    included_classes = y.unique(return_counts=True)[1] >= threshold\n",
        "    filter = included_classes[y]\n",
        "    # remap labels (i.e. if they were 0-8 and we remove 4 labels, new labels should be between 0 and 4)\n",
        "    label_mapping = torch.zeros(included_classes.size(dim=0), dtype=torch.long)\n",
        "    label_mapping[included_classes] = torch.arange(torch.count_nonzero(included_classes))\n",
        "\n",
        "    y = label_mapping[y[filter]].to(torch.long)\n",
        "    X = X[filter]\n",
        "\n",
        "    # remove edges that had their nodes removed\n",
        "    edge_index = filter_edge_index(edge_index, filter)\n",
        "\n",
        "    return Data(x=X, y=y, edge_index=edge_index)\n",
        "\n",
        "def get_adjacency_matrix(dataset):\n",
        "    edge_index = dataset.edge_index\n",
        "\n",
        "    # make sparse adjacency matrix, A\n",
        "    values = torch.ones(edge_index.size(dim=1), dtype = torch.int)\n",
        "    A = torch.sparse_coo_tensor(edge_index, values, (dataset.x.size(dim=0), dataset.x.size(dim=0)), dtype=torch.float)\n",
        "\n",
        "    return A"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aViLSsl2VD9K",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "79ef62f7-9269-4edc-e9cb-ed5393690ec3"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Downloading https://github.com/shchur/gnn-benchmark/raw/master/data/npz/amazon_electronics_computers.npz\n",
            "Processing...\n",
            "Done!\n"
          ]
        }
      ],
      "source": [
        "from torch_geometric.datasets import Amazon\n",
        "from torch_geometric.loader import DataLoader\n",
        "\n",
        "dataset = Amazon('.', name='Computers')[0]\n",
        "# prepare dataset by removing classes that have less than 1000 examples\n",
        "dataset = prepare_dataset(dataset, 1000)\n",
        "# get num classes\n",
        "num_classes = torch.unique(dataset.y).size(dim=0)\n",
        "\n",
        "# train/test split\n",
        "train_dataset, test_dataset = train_test_split(dataset, 0.2)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "edge_index = torch.tensor([[0, 2, 0, 0, 2, 3, 1, 4, 1, 2, 4],\n",
        "                           [1, 0, 3, 4, 1, 1, 4, 4, 1, 3, 3]], dtype=torch.long)\n",
        "x = torch.tensor([[0], [1], [2], [3], [4]], dtype=torch.float)\n",
        "y = torch.tensor([0, 1, 1, 2, 2], dtype=torch.long)\n",
        "\n",
        "data = Data(x=x, y=y, edge_index=edge_index)\n",
        "\n",
        "# filter_edge_index(data.edge_index, torch.tensor([False, True, True, False, True]))\n",
        "\n",
        "# train_data, test_data = train_test_split(data, 0)\n",
        "\n",
        "data = prepare_dataset(data, 1)\n",
        "A = get_adjacency_matrix(data)\n",
        "\n",
        "torch.sparse.sum(torch.index_select(A, 0, torch.tensor([0, 1, 2, 3])))\n",
        "\n",
        "# .to_dense()\n",
        "# print(data.y)\n",
        "# print(data.x)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "W_RHNlMG9Vkg",
        "outputId": "ca3613d8-7640-4650-b7fd-19be5695953a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor(9.)"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "29Swawd9we8f"
      },
      "outputs": [],
      "source": [
        "X_train, y_train, X_test, y_test = train_dataset.x, train_dataset.y, test_dataset.x, test_dataset.y\n",
        "A_train, A_test = get_adjacency_matrix(train_dataset), get_adjacency_matrix(test_dataset)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a4Md2Ds3q3CY"
      },
      "source": [
        "## Train/Test\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nu0AtfY7qzba"
      },
      "outputs": [],
      "source": [
        "# train\n",
        "def train(X, y, model, loss_fn, optimizer): \n",
        "    # make this into dataloader using backup\n",
        "    model.train()\n",
        "    X, y = X.to(device), y.to(device)\n",
        "\n",
        "    # Compute prediction error\n",
        "    pred = model(X)\n",
        "    loss = loss_fn(pred, y)\n",
        "\n",
        "    # Backpropagation\n",
        "    optimizer.zero_grad()\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "# test\n",
        "def test(X, y, split, model, loss_fn):\n",
        "    size = X.size(dim=0)\n",
        "    model.eval()\n",
        "    test_loss, correct = 0, 0\n",
        "    with torch.inference_mode():\n",
        "        X, y = X, y\n",
        "        X, y = X.to(device), y.to(device)\n",
        "        pred = model(X)\n",
        "        test_loss += loss_fn(pred, y).item()\n",
        "        correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
        "    correct /= size\n",
        "    print(f\"{split.title()} Error: \\n Accuracy: {(100*correct):>0.1f}%, Loss: {test_loss:>8f} \\n\")\n",
        "\n",
        "# graph train\n",
        "def graph_train(X, y, A, model, loss_fn, optimizer): \n",
        "    # make this into dataloader using backup\n",
        "    model.train()\n",
        "    model.set_adj_mat(A)\n",
        "    X, y = X.to(device), y.to(device)\n",
        "\n",
        "    # Compute prediction error\n",
        "    pred = model(X)\n",
        "    loss = loss_fn(pred, y)\n",
        "\n",
        "    # Backpropagation\n",
        "    optimizer.zero_grad()\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "# graph test\n",
        "def graph_test(X, y, A, split, model, loss_fn):\n",
        "    size = X.size(dim=0)\n",
        "    model.set_adj_mat(A)\n",
        "    model.eval()\n",
        "    test_loss, correct = 0, 0\n",
        "    with torch.inference_mode():\n",
        "        X, y = X, y\n",
        "        X, y = X.to(device), y.to(device)\n",
        "        pred = model(X)\n",
        "        test_loss += loss_fn(pred, y).item()\n",
        "        correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
        "    correct /= size\n",
        "    print(f\"{split.title()} Error: \\n Accuracy: {(100*correct):>0.1f}%, Loss: {test_loss:>8f} \\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gdmoD-13Ad2C"
      },
      "source": [
        "##Sampling from K-hop neighbourhood"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0CGK4KdXAdVz"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SFvQEYa3j3O7"
      },
      "source": [
        "## Encoder"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XixvkAQpqcFM"
      },
      "source": [
        "Encoder Design\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9ZT6RQbkiTbj"
      },
      "outputs": [],
      "source": [
        "# encoder\n",
        "dimensions = [767, 300, 60]\n",
        "encoder_train = nn.Sequential(\n",
        "    MLP(dimensions),\n",
        "    nn.Linear(dimensions[-1], num_classes),\n",
        "    nn.Softmax(dim=1)\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B8Jk2-tbqLer"
      },
      "source": [
        "Encoder Pretraining"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0HYGNjZgqo7k",
        "outputId": "99ce441b-8585-4d85-8a54-53c7dcf77c5c"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train Error: \n",
            " Accuracy: 91.4%, Loss: 0.835441 \n",
            "\n",
            "Train Error: \n",
            " Accuracy: 93.5%, Loss: 0.813172 \n",
            "\n",
            "Train Error: \n",
            " Accuracy: 94.8%, Loss: 0.801979 \n",
            "\n",
            "Train Error: \n",
            " Accuracy: 95.1%, Loss: 0.799237 \n",
            "\n",
            "Train Error: \n",
            " Accuracy: 95.3%, Loss: 0.795687 \n",
            "\n",
            "Train Error: \n",
            " Accuracy: 95.5%, Loss: 0.792990 \n",
            "\n",
            "Train Error: \n",
            " Accuracy: 95.7%, Loss: 0.790494 \n",
            "\n",
            "Train Error: \n",
            " Accuracy: 95.9%, Loss: 0.788637 \n",
            "\n",
            "Train Error: \n",
            " Accuracy: 96.0%, Loss: 0.787141 \n",
            "\n",
            "Train Error: \n",
            " Accuracy: 96.1%, Loss: 0.785771 \n",
            "\n",
            "Test Error: \n",
            " Accuracy: 91.4%, Loss: 0.830850 \n",
            "\n",
            "Done!\n"
          ]
        }
      ],
      "source": [
        "encoder_model = encoder_train.to(device)\n",
        "loss_fn = nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.Adam(encoder_model.parameters(), lr=1e-3)\n",
        "\n",
        "# if node_level:\n",
        "#   optimizer = op.optimizers.optimizer.DPOptimizer(\n",
        "#       # TODO: Fill out these parameters '?'\n",
        "#       optimizer=optimizer,\n",
        "#       noise_multiplier=?,\n",
        "#       max_grad_norm=?\n",
        "#   )\n",
        "\n",
        "epochs = 100\n",
        "for t in range(epochs):\n",
        "    train(X_train, y_train, encoder_model, loss_fn, optimizer)\n",
        "    if (t + 1) % 10 == 0:\n",
        "      test(X_train, y_train, \"TRAIN\", encoder_model, loss_fn)\n",
        "test(X_test, y_test, \"TEST\", encoder_model, loss_fn)\n",
        "print(\"Done!\")\n",
        "\n",
        "encoder = encoder_model[0]\n",
        "\n",
        "# for name, param in encoder_model.named_parameters():\n",
        "#     if param.requires_grad:\n",
        "#         print(name, param.data)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Full Model Training"
      ],
      "metadata": {
        "id": "Zqf6_kLMUZUn"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "TODO:\n",
        "  - try other dataset\n",
        "  - compare with their implementation\n",
        "  - try train/test split across edges\n",
        "  - batch normalization\n",
        "  - add identity matrix to see performance"
      ],
      "metadata": {
        "id": "2os_fvuNJmcJ"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lo_SC_jUkl-y"
      },
      "source": [
        "Train full model"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "encoder.requires_grad=False\n",
        "# model = nn.Sequential(encoder, \n",
        "#                       PMA(K_hop, sigma), \n",
        "#                       nn.Linear(60, num_classes), \n",
        "#                       nn.Softmax(dim=1))\n",
        "# model = nn.Sequential(encoder,\n",
        "#                       PMA(K_hop, sigma),\n",
        "#                       Classification(K_hop, [60, 20], [(K_hop+1)*20, num_classes]))\n",
        "model = GAP(encoder, \n",
        "            PMA(K_hop, sigma), \n",
        "            Classification(K_hop, [60, 20], [(K_hop+1)*20, num_classes]))\n",
        "# model = nn.Sequential(encoder, \n",
        "#                       PMA(K_hop, sigma))\n",
        "model = model.to(device)\n",
        "loss_fn = nn.CrossEntropyLoss()\n",
        "\n",
        "# print(model(X))"
      ],
      "metadata": {
        "id": "UPx__jwgLUCL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "optimizer = torch.optim.SGD(model.parameters(), lr=1e-1)\n",
        "\n",
        "epochs = 10\n",
        "# model.set_adj_mat(A_train)\n",
        "for t in range(epochs):\n",
        "    graph_train(X_train, y_train, A_train, model, loss_fn, optimizer)\n",
        "    if (t + 1) % 2 == 0:\n",
        "      graph_test(X_train, y_train, A_train, \"TRAIN\", model, loss_fn)\n",
        "# model.set_adj_mat(A_test)\n",
        "graph_test(X_test, y_test, A_test, \"TEST\", model, loss_fn)\n",
        "print(\"Done!\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "z22WhxLcaveo",
        "outputId": "636a2254-6361-4c79-c977-6d8e7288f8ec"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train Error: \n",
            " Accuracy: 95.6%, Loss: 0.160033 \n",
            "\n",
            "Train Error: \n",
            " Accuracy: 95.6%, Loss: 0.159795 \n",
            "\n",
            "Train Error: \n",
            " Accuracy: 95.6%, Loss: 0.159560 \n",
            "\n",
            "Train Error: \n",
            " Accuracy: 95.6%, Loss: 0.159328 \n",
            "\n",
            "Train Error: \n",
            " Accuracy: 95.6%, Loss: 0.159099 \n",
            "\n",
            "Test Error: \n",
            " Accuracy: 91.4%, Loss: 0.256538 \n",
            "\n",
            "Done!\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x0yD8RqskkhV",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 389
        },
        "outputId": "b9900a7b-b1ed-400a-de66-3d9a1ee3841e"
      },
      "outputs": [
        {
          "output_type": "error",
          "ename": "IndexError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-276-55001f8cf98d>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     17\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mt\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mepochs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m     \u001b[0;31m# print(f\"Epoch {t+1}\\n-------------------------------\")\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 19\u001b[0;31m     \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgap_model\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mloss_fn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     20\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mt\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0;36m10\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m       \u001b[0mtest\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgap_model\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mloss_fn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-225-d40dcccb3455>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(X, y, model, loss_fn, optimizer)\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m     \u001b[0;31m# Compute prediction error\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 8\u001b[0;31m     \u001b[0mpred\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      9\u001b[0m     \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mloss_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpred\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     10\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1128\u001b[0m         \u001b[0mThe\u001b[0m \u001b[0mhook\u001b[0m \u001b[0mcan\u001b[0m \u001b[0mmodify\u001b[0m \u001b[0mthe\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m.\u001b[0m \u001b[0mUser\u001b[0m \u001b[0mcan\u001b[0m \u001b[0meither\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0ma\u001b[0m \u001b[0mtuple\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0ma\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1129\u001b[0m         \u001b[0msingle\u001b[0m \u001b[0mmodified\u001b[0m \u001b[0mvalue\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mthe\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m.\u001b[0m \u001b[0mWe\u001b[0m \u001b[0mwill\u001b[0m \u001b[0mwrap\u001b[0m \u001b[0mthe\u001b[0m \u001b[0mvalue\u001b[0m \u001b[0minto\u001b[0m \u001b[0ma\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1130\u001b[0;31m         \u001b[0;32mif\u001b[0m \u001b[0ma\u001b[0m \u001b[0msingle\u001b[0m \u001b[0mvalue\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0mreturned\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0munless\u001b[0m \u001b[0mthat\u001b[0m \u001b[0mvalue\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0malready\u001b[0m \u001b[0ma\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1131\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1132\u001b[0m         \u001b[0mReturns\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-8-c8ae150ef23a>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     16\u001b[0m     \u001b[0mcache\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpma\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx_encoded\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m     \u001b[0;31m# classification\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 18\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclassification\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcache\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1128\u001b[0m         \u001b[0mThe\u001b[0m \u001b[0mhook\u001b[0m \u001b[0mcan\u001b[0m \u001b[0mmodify\u001b[0m \u001b[0mthe\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m.\u001b[0m \u001b[0mUser\u001b[0m \u001b[0mcan\u001b[0m \u001b[0meither\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0ma\u001b[0m \u001b[0mtuple\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0ma\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1129\u001b[0m         \u001b[0msingle\u001b[0m \u001b[0mmodified\u001b[0m \u001b[0mvalue\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mthe\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m.\u001b[0m \u001b[0mWe\u001b[0m \u001b[0mwill\u001b[0m \u001b[0mwrap\u001b[0m \u001b[0mthe\u001b[0m \u001b[0mvalue\u001b[0m \u001b[0minto\u001b[0m \u001b[0ma\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1130\u001b[0;31m         \u001b[0;32mif\u001b[0m \u001b[0ma\u001b[0m \u001b[0msingle\u001b[0m \u001b[0mvalue\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0mreturned\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0munless\u001b[0m \u001b[0mthat\u001b[0m \u001b[0mvalue\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0malready\u001b[0m \u001b[0ma\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1131\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1132\u001b[0m         \u001b[0mReturns\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-113-9e3d826c6d70>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, cache)\u001b[0m\n\u001b[1;32m     14\u001b[0m         \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbase_mlps\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 16\u001b[0;31m           \u001b[0mencoding\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbase_mlps\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcache\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     17\u001b[0m           \u001b[0mout\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mencoding\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# add corresponding encoding\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m           \u001b[0;31m# TEMP\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mIndexError\u001b[0m: too many indices for tensor of dimension 2"
          ]
        }
      ],
      "source": [
        "# Sigma calculated above in node-level and edge-level DP case\n",
        "gap = GAP(encoder, PMA(A, K_hop, sigma), Classification(K_hop, [60, 30, 20], [(K_hop+1)*20, 60, 30, num_classes]))\n",
        "gap_model = gap.to(device)\n",
        "loss_fn = nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.SGD(gap_model.parameters(), lr=1e-1)\n",
        "\n",
        "# if node_level:\n",
        "#   optimizer = op.optimizers.optimizer.DPOptimizer(\n",
        "#       # TODO: Fill out these parameters '?'\n",
        "#       optimizer=optimizer,\n",
        "#       noise_multiplier=?,\n",
        "#       max_grad_norm=?,\n",
        "#       loss_reduction='sum'\n",
        "#   ) \n",
        "\n",
        "epochs = 500\n",
        "for t in range(epochs):\n",
        "    # print(f\"Epoch {t+1}\\n-------------------------------\")\n",
        "    train(X, y, gap_model, loss_fn, optimizer)\n",
        "    if t % 10 == 0:\n",
        "      test(X, y, gap_model, loss_fn)\n",
        "print(\"Done!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BgX93DaHqqzh"
      },
      "source": [
        "## Backup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vof5FsrKqroZ"
      },
      "outputs": [],
      "source": [
        "# # train\n",
        "# def train(dataloader, model, loss_fn, optimizer, print_every = 100):\n",
        "#     size = len(dataloader.dataset)\n",
        "#     model.train()\n",
        "#     for batch, (X, y) in enumerate(dataloader):\n",
        "#         X, y = X.to(device), y.to(device)\n",
        "\n",
        "#         # Compute prediction error\n",
        "#         pred = model(X)\n",
        "#         loss = loss_fn(pred, y)\n",
        "\n",
        "#         # Backpropagation\n",
        "#         optimizer.zero_grad()\n",
        "#         loss.backward()\n",
        "#         optimizer.step()\n",
        "\n",
        "#         if batch % print_every == 0:\n",
        "#             loss, current = loss.item(), batch * len(X)\n",
        "#             print(f\"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]\")\n",
        "\n",
        "# # test\n",
        "# def test(dataloader, model, loss_fn):\n",
        "#     size = len(dataloader.dataset)\n",
        "#     num_batches = len(dataloader)\n",
        "#     model.eval()\n",
        "#     test_loss, correct = 0, 0\n",
        "#     with torch.inference_mode():\n",
        "#         for X, y in dataloader:\n",
        "#             X, y = X.to(device), y.to(device)\n",
        "#             pred = model(X)\n",
        "#             test_loss += loss_fn(pred, y).item()\n",
        "#             correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
        "#     test_loss /= num_batches\n",
        "#     correct /= size\n",
        "#     print(f\"Test Error: \\n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \\n\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "vb1hVZklFoDP",
        "Ty_dao27oEAB",
        "WebCohAOIX0P",
        "a4Md2Ds3q3CY",
        "gdmoD-13Ad2C",
        "BgX93DaHqqzh"
      ],
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}