ModernGAP_PMAT_Implementation.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_PMAT_Implementation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aFQCNsDWU5mf",
        "outputId": "1c41189d-02c4-40df-9dd5-494db0c314eb"
      },
      "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.1.0%2Bpt112cpu-cp38-cp38-linux_x86_64.whl (315 kB)\n",
            "\u001b[K     |████████████████████████████████| 315 kB 4.2 MB/s \n",
            "\u001b[?25hCollecting torch-sparse\n",
            "  Downloading https://data.pyg.org/whl/torch-1.12.0%2Bcpu/torch_sparse-0.6.15%2Bpt112cpu-cp38-cp38-linux_x86_64.whl (650 kB)\n",
            "\u001b[K     |████████████████████████████████| 650 kB 23.7 MB/s \n",
            "\u001b[?25hCollecting torch-cluster\n",
            "  Downloading https://data.pyg.org/whl/torch-1.12.0%2Bcpu/torch_cluster-1.6.0%2Bpt112cpu-cp38-cp38-linux_x86_64.whl (324 kB)\n",
            "\u001b[K     |████████████████████████████████| 324 kB 293 kB/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%2Bpt112cpu-cp38-cp38-linux_x86_64.whl (128 kB)\n",
            "\u001b[K     |████████████████████████████████| 128 kB 321 kB/s \n",
            "\u001b[?25hCollecting torch-geometric==2.0.0\n",
            "  Downloading torch_geometric-2.0.0.tar.gz (308 kB)\n",
            "\u001b[K     |████████████████████████████████| 308 kB 15.0 MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.8/dist-packages (from torch-geometric==2.0.0) (1.21.6)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.8/dist-packages (from torch-geometric==2.0.0) (4.64.1)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.8/dist-packages (from torch-geometric==2.0.0) (1.7.3)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.8/dist-packages (from torch-geometric==2.0.0) (2.6.3)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.8/dist-packages (from torch-geometric==2.0.0) (1.0.2)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.8/dist-packages (from torch-geometric==2.0.0) (2.23.0)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.8/dist-packages (from torch-geometric==2.0.0) (1.3.5)\n",
            "Collecting rdflib\n",
            "  Downloading rdflib-6.2.0-py3-none-any.whl (500 kB)\n",
            "\u001b[K     |████████████████████████████████| 500 kB 40.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: googledrivedownloader in /usr/local/lib/python3.8/dist-packages (from torch-geometric==2.0.0) (0.4)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.8/dist-packages (from torch-geometric==2.0.0) (2.11.3)\n",
            "Requirement already satisfied: pyparsing in /usr/local/lib/python3.8/dist-packages (from torch-geometric==2.0.0) (3.0.9)\n",
            "Collecting yacs\n",
            "  Downloading yacs-0.1.8-py3-none-any.whl (14 kB)\n",
            "Requirement already satisfied: PyYAML in /usr/local/lib/python3.8/dist-packages (from torch-geometric==2.0.0) (6.0)\n",
            "Requirement already satisfied: MarkupSafe>=0.23 in /usr/local/lib/python3.8/dist-packages (from jinja2->torch-geometric==2.0.0) (2.0.1)\n",
            "Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.8/dist-packages (from pandas->torch-geometric==2.0.0) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.8/dist-packages (from pandas->torch-geometric==2.0.0) (2022.6)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.8/dist-packages (from python-dateutil>=2.7.3->pandas->torch-geometric==2.0.0) (1.15.0)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.8/dist-packages (from rdflib->torch-geometric==2.0.0) (57.4.0)\n",
            "Collecting isodate\n",
            "  Downloading isodate-0.6.1-py2.py3-none-any.whl (41 kB)\n",
            "\u001b[K     |████████████████████████████████| 41 kB 289 kB/s \n",
            "\u001b[?25hRequirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests->torch-geometric==2.0.0) (2.10)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests->torch-geometric==2.0.0) (3.0.4)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests->torch-geometric==2.0.0) (1.24.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests->torch-geometric==2.0.0) (2022.9.24)\n",
            "Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.8/dist-packages (from scikit-learn->torch-geometric==2.0.0) (1.2.0)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from scikit-learn->torch-geometric==2.0.0) (3.1.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.0.0-py3-none-any.whl size=513272 sha256=7c9c3d70d7ff46a49b7d6c13f07a5278fb5ef09a4136f6072de1ce4f5a4e5780\n",
            "  Stored in directory: /root/.cache/pip/wheels/6c/be/4c/bd8d186bc6ed4a161041903706bb1e99b74dc560982c23718e\n",
            "Successfully built torch-geometric\n",
            "Installing collected packages: isodate, yacs, rdflib, torch-spline-conv, torch-sparse, torch-scatter, torch-geometric, torch-cluster\n",
            "Successfully installed isodate-0.6.1 rdflib-6.2.0 torch-cluster-1.6.0+pt112cpu torch-geometric-2.0.0 torch-scatter-2.1.0+pt112cpu torch-sparse-0.6.15+pt112cpu torch-spline-conv-1.2.1+pt112cpu yacs-0.1.8\n",
            "\u001b[33mWARNING: Skipping pyvacy as it is not installed.\u001b[0m\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting pyvacy\n",
            "  Downloading pyvacy-0.0.32.tar.gz (10 kB)\n",
            "Building wheels for collected packages: pyvacy\n",
            "  Building wheel for pyvacy (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pyvacy: filename=pyvacy-0.0.32-py3-none-any.whl size=12401 sha256=e08c611b1d59d4ad4a6d82bf249e6610a61abf5b0942d414b0caf29390413310\n",
            "  Stored in directory: /root/.cache/pip/wheels/0f/b9/fd/45fb9afde4d0efd0022f976cf3e53f058e01cf99812687e36f\n",
            "Successfully built pyvacy\n",
            "Installing collected packages: pyvacy\n",
            "Successfully installed pyvacy-0.0.32\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "if torch.cuda.is_available():\n",
        "  #NVIDIA GPU version\n",
        "  %pip uninstall torch-scatter torch-sparse torch-geometric torch-cluster  --y\n",
        "  %pip install torch-scatter -f https://data.pyg.org/whl/torch-{torch.__version__}.html\n",
        "  %pip install torch-sparse -f https://data.pyg.org/whl/torch-{torch.__version__}.html\n",
        "  %pip install torch-cluster -f https://data.pyg.org/whl/torch-{torch.__version__}.html\n",
        "  %pip install git+https://github.com/pyg-team/pytorch_geometric.git\n",
        "else:\n",
        "  #CPU version\n",
        "  %pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric==2.0.0 -f https://data.pyg.org/whl/torch-1.12.0+cpu.html\n",
        "%pip uninstall pyvacy  --y\n",
        "%pip install pyvacy\n",
        "\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vb1hVZklFoDP"
      },
      "source": [
        "## Encoder Module"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "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_selu_stack = nn.Sequential(*layers)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.flatten(x)\n",
        "        logits = self.linear_selu_stack(x)\n",
        "        return logits"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GF264FWv0pJI"
      },
      "source": [
        "## PMA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "OMzLWqwF03Uu"
      },
      "outputs": [],
      "source": [
        "class AggregationModule(nn.Module):\n",
        "  edge_index = None\n",
        "\n",
        "  def __init__(self):\n",
        "        super().__init__()\n",
        "\n",
        "class PMA(AggregationModule):\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",
        "        self.num_hops = num_hops\n",
        "        self.sigma = sigma\n",
        "\n",
        "    def forward(self, x):\n",
        "        # TEMP SOLUTION\n",
        "        if AggregationModule.edge_index is None:\n",
        "          raise RuntimeError(\"Set AggregationModule.edge_index [TEMP SOLUTION] before running\")\n",
        "        edge_index = AggregationModule.edge_index\n",
        "        A = get_adjacency_matrix(edge_index, x.size(dim=0))\n",
        "        out = [torch.nn.functional.normalize(x, dim=1)]\n",
        "        for k in range(self.num_hops):\n",
        "            aggr = torch.mm(torch.transpose(A, 0, 1), 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)\n",
        "\n",
        "class PMAT(AggregationModule):\n",
        "    # def __init__(self, num_hops, transform_dimensions):\n",
        "    def __init__(self, num_hops, encoding_dimensions, sigma):\n",
        "        super().__init__()\n",
        "        self.num_hops = num_hops\n",
        "        self.sigma = sigma\n",
        "        self.sigmoid = nn.Sigmoid()\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([2*transform_dimensions[-1], 1])) # Attention mechanism takes 2 encodings and outputs 1 weight\n",
        "          # TODO: Figure out if we want a transformer?\n",
        "          self.attentions.append(MLP([2*encoding_dimensions, 1]))\n",
        "\n",
        "    def forward(self, x):\n",
        "        # TEMP SOLUTION\n",
        "        if AggregationModule.edge_index is None:\n",
        "          raise RuntimeError(\"Set AggregationModule.edge_index [TEMP SOLUTION] before running\")\n",
        "        edge_index = AggregationModule.edge_index\n",
        "        out = [torch.nn.functional.normalize(x, dim=1)]\n",
        "        for k in range(self.num_hops):\n",
        "            # Do we need to do a transform? I reckon we can use raw encoding and aggregate according to attention (and then the Classification module)\n",
        "            # can handle how the aggregations get transformed\n",
        "            # h = self.transforms[k](out[-1])\n",
        "            h = out[-1]\n",
        "            e_values = self.attentions[k](h[edge_index.T].reshape(edge_index.size(dim=1), 2*h.size(dim=1))) # DPSGD to guarantee DP attention training\n",
        "            # we have to use Sigmoid because if we use Softmax, removing an edge will change the weight of all other edges in the neighbourhood\n",
        "            alpha_values = self.sigmoid(e_values)\n",
        "            alpha = torch.sparse_coo_tensor(edge_index,\n",
        "                                            alpha_values.reshape(edge_index.size(dim=1)),\n",
        "                                            (x.size(dim=0), x.size(dim=0)),\n",
        "                                            dtype=torch.float).transpose(0, 1)\n",
        "\n",
        "            aggr = torch.sparse.mm(alpha, h)\n",
        "            # Might need to not use \"transforms\" and instead do raw aggregations like the original PMA\n",
        "            # aggr = torch.mm(self.A, out[-1]) \n",
        "            # noised = aggr # TODO: add noise # Gaussian mechanism to guarantee DP for neighbourhood aggregation\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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "CZnwBQti_gK-",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b604f503-6622-4f6e-f0ce-eade22592966"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([[[0.0000, 1.0000, 0.0000],\n",
            "         [0.3333, 0.6667, 0.6667],\n",
            "         [0.5345, 0.8018, 0.2673],\n",
            "         [0.5571, 0.3714, 0.7428],\n",
            "         [0.7428, 0.3714, 0.5571]],\n",
            "\n",
            "        [[0.5345, 0.8018, 0.2673],\n",
            "         [0.3986, 0.7821, 0.4790],\n",
            "         [0.0000, 0.0000, 0.0000],\n",
            "         [0.4821, 0.8186, 0.3121],\n",
            "         [0.4117, 0.7761, 0.4777]],\n",
            "\n",
            "        [[0.0000, 0.0000, 0.0000],\n",
            "         [0.4729, 0.8053, 0.3575],\n",
            "         [0.0000, 0.0000, 0.0000],\n",
            "         [0.4748, 0.7945, 0.3785],\n",
            "         [0.4492, 0.7916, 0.4141]]], grad_fn=<StackBackward0>)\n"
          ]
        }
      ],
      "source": [
        "# TEMP CODE\n",
        "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, 0], [1, 2, 2], [2, 3, 1], [3, 2, 4], [4, 2, 3]], dtype=torch.float)\n",
        "y = torch.tensor([0, 1, 1, 2, 2], dtype=torch.long)\n",
        "pma = PMAT(2, 3, 0)\n",
        "AggregationModule.edge_index = edge_index\n",
        "stack = pma(x)\n",
        "print(stack)\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": 7,
      "metadata": {
        "id": "Oaw4OM5_2Tmk",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f98d91d7-0f8a-4300-8742-f0aabcd5e1df"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[ 9.6033, 10.6737, 10.8559],\n",
              "        [10.6253,  9.8901,  8.9841]])"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ],
      "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": 8,
      "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",
        "        self.num_hops = num_hops\n",
        "        if encoder_dimensions:\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(self.num_hops+1):\n",
        "          if self.base_mlps:\n",
        "            encoding = self.base_mlps[i](cache[i,:,:])\n",
        "            out.append(encoding) # add corresponding encoding\n",
        "          else:\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": 9,
      "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 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": 10,
      "metadata": {
        "id": "G_Pf4P_bIZ5S",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0c5fdf39-c80c-4e53-8238-b9a675374130"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epsilon: 1.00, Delta: 0.50, Sigma: 2.61\n"
          ]
        }
      ],
      "source": [
        "node_level = True\n",
        "\n",
        "# Edge level DP\n",
        "agg_epsilon, agg_delta, epsilon_1, epsilon_5, pmat_epsilon, pmat_delta = 1, 0.5, 4, 4, 4, 0.5\n",
        "K_hop = 3\n",
        "agg_sigma = 1 / np.max(np.roots([K_hop/2, np.sqrt(2*K_hop*np.log(1/agg_delta)), -agg_epsilon]))\n",
        "# Node level DP\n",
        "if (node_level):\n",
        "  pass\n",
        "  # How do we calculate this?\n",
        "data = \"reddit\"\n",
        "\n",
        "print(f\"Epsilon: {agg_epsilon:>0.2f}, Delta: {agg_delta:>0.2f}, Sigma: {agg_sigma:>0.2f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Preprocess"
      ],
      "metadata": {
        "id": "XDJP3FhOi_Ty"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "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 add_edge_to_low_degree_nodes(dataset, low_degree_threshold):\n",
        "    X, y, edge_index = dataset.x, dataset.y, dataset.edge_index\n",
        "\n",
        "    # get low degree nodes\n",
        "    A = get_adjacency_matrix(edge_index, X.size(dim=0))\n",
        "    sums = torch.sparse.sum(A, dim=1).to_dense()\n",
        "    mask = sums < low_degree_threshold\n",
        "    \n",
        "    # get edge_index mask for neighbours of low degree nodes\n",
        "    filter = mask[edge_index[0, :]]\n",
        "    low_degree_edges_index = edge_index[:, filter]\n",
        "    low_degree_A = get_adjacency_matrix(low_degree_edges_index, X.size(dim=0))\n",
        "\n",
        "    # get 1-hop neighbours and add to A \n",
        "    # NOTE: without sampling (just adds all 1-hop neighbours)\n",
        "    one_hop_low_degree_A = torch.sparse.mm(low_degree_A, A)\n",
        "    new_edge_index = torch.add(one_hop_low_degree_A, A).coalesce().indices()\n",
        "    return Data(x=X, y=y, edge_index=new_edge_index)\n",
        "    \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",
        "# make sparse adjacency matrix, A\n",
        "def get_adjacency_matrix(edge_index, num_nodes):\n",
        "    values = torch.ones(edge_index.size(dim=1), dtype = torch.int)\n",
        "    A = torch.sparse_coo_tensor(edge_index, values, (num_nodes, num_nodes), dtype=torch.float)\n",
        "    return A\n",
        "\n",
        "def standardization(train_dataset, test_dataset):\n",
        "    X = train_dataset.x\n",
        "    means = X.mean(dim=0, keepdim=True)\n",
        "    stds = X.std(dim=0, keepdim=True)\n",
        "    X_train = (X - means) / stds\n",
        "    X_test = (test_dataset.x - means) / stds\n",
        "    return Data(x=X_train, y=train_dataset.y, edge_index=train_dataset.edge_index), Data(x=X_test, y=test_dataset.y, edge_index=test_dataset.edge_index)\n",
        "\n",
        "def add_self_edges(dataset):\n",
        "    X = dataset.x\n",
        "    self_edges = torch.stack((torch.arange(X.size(dim=0)), torch.arange(X.size(dim=0))))\n",
        "    edge_index = torch.cat((dataset.edge_index, self_edges), dim=1)\n",
        "    return Data(x=X, y=dataset.y, edge_index=edge_index)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E7oEXJ63oyl-"
      },
      "source": [
        "## Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "aViLSsl2VD9K"
      },
      "outputs": [],
      "source": [
        "# from torch_geometric.datasets import Amazon\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",
        "# dataset = add_self_edges(dataset)\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)\n",
        "# print(train_dataset.edge_index.size(dim=1))\n",
        "# train_dataset = add_edge_to_low_degree_nodes(train_dataset, 10)\n",
        "# print(train_dataset.edge_index.size(dim=1))\n",
        "# test_dataset = add_edge_to_low_degree_nodes(test_dataset, 10)\n",
        "# train_dataset, test_dataset = standardization(train_dataset, test_dataset)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from torch_geometric.datasets import Reddit\n",
        "dataset = Reddit('.')[0]\n",
        "# prepare dataset by removing classes that have less than 1000 examples\n",
        "dataset = prepare_dataset(dataset, 10000)\n",
        "dataset = add_self_edges(dataset)\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)\n",
        "print(train_dataset.edge_index.size(dim=1))\n",
        "# train_dataset = add_edge_to_low_degree_nodes(train_dataset, 10)\n",
        "print(train_dataset.edge_index.size(dim=1))\n",
        "# test_dataset = add_edge_to_low_degree_nodes(test_dataset, 10)\n",
        "train_dataset, test_dataset = standardization(train_dataset, test_dataset)"
      ],
      "metadata": {
        "id": "JrAKSDBl03mk",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d6dddc31-e267-42db-9648-87a78335d8c8"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Downloading https://data.dgl.ai/dataset/reddit.zip\n",
            "Extracting ./raw/reddit.zip\n",
            "Processing...\n",
            "Done!\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "29524966\n",
            "29524966\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "29Swawd9we8f",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a8faef98-d0c6-4486-ae53-472f2746052f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([93592, 602])\n"
          ]
        }
      ],
      "source": [
        "from torch_geometric.loader import NeighborLoader\n",
        "\n",
        "X_train, y_train, edge_index_train = train_dataset.x, train_dataset.y, train_dataset.edge_index\n",
        "X_test, y_test, edge_index_test = test_dataset.x, test_dataset.y, test_dataset.edge_index\n",
        "\n",
        "# using large number like 10,000 so that all neighbours are sampled \n",
        "# I don't like how it samples, so I'm just gonna sample everything\n",
        "train_loader = NeighborLoader(train_dataset, num_neighbors=[X_train.size(dim=0)]*K_hop, batch_size=X_train.size(dim=0))\n",
        "print(X_train.size())"
      ]
    },
    {
      "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, 0], [1, 2, 2], [2, 3, 1], [3, 2, 4], [4, 2, 3]], dtype=torch.float)\n",
        "# y = torch.tensor([0, 1, 1, 2, 2], dtype=torch.long)\n",
        "# data = Data(x=x, y=y, edge_index=edge_index)\n",
        "# add_self_edges(data)\n",
        "# data = add_edge_to_low_degree_nodes(data, 3)\n",
        "# print(data.x)\n",
        "# data = standardization(data)\n",
        "# print(data.x)\n",
        "# print(data.x.std(dim=0))\n",
        "# print(data.edge_index)\n",
        "# print(get_adjacency_matrix(data.edge_index, x.size(dim=0)).to_dense())\n",
        "# train, test = train_test_split(data, 0.2)\n",
        "# print(train.x)\n",
        "# print(train.edge_index)\n",
        "# print(test.x)\n",
        "# print(test.edge_index)"
      ],
      "metadata": {
        "id": "Ot9FV6_9Ok3e"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a4Md2Ds3q3CY"
      },
      "source": [
        "## Train/Test\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "nu0AtfY7qzba"
      },
      "outputs": [],
      "source": [
        "# train\n",
        "def train(batch, model, loss_fn, optimizer):\n",
        "  model.train()\n",
        "  batch = batch.to(device)\n",
        "  AggregationModule.edge_index = batch.edge_index\n",
        "  # compute prediction error\n",
        "  pred = model(batch)\n",
        "  loss = loss_fn(pred, batch.y)\n",
        "  # backpropagation\n",
        "  optimizer.zero_grad()\n",
        "  loss.backward()\n",
        "  optimizer.step()\n",
        "\n",
        "# test\n",
        "def test(X, y, edge_index, 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.to(device), y.to(device)\n",
        "        AggregationModule.edge_index = edge_index.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}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SFvQEYa3j3O7"
      },
      "source": [
        "## Encoder"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XixvkAQpqcFM"
      },
      "source": [
        "Encoder Design\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "9ZT6RQbkiTbj"
      },
      "outputs": [],
      "source": [
        "# encoder\n",
        "dimensions = [602, 300, 60]\n",
        "encoder_model = 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": 18,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 550
        },
        "id": "0HYGNjZgqo7k",
        "outputId": "85825b94-9396-40e3-95b3-98b0abe65fef"
      },
      "outputs": [
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-18-12c840e063cc>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[0mepochs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m80\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     14\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[0;32m---> 15\u001b[0;31m     \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mencoder_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     16\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mt\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\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     17\u001b[0m       \u001b[0mtest\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX_train\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_train\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_index_train\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"TRAIN\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mencoder_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-16-f7ce8442880d>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(dataloader, model, loss_fn, optimizer)\u001b[0m\n\u001b[1;32m     14\u001b[0m         \u001b[0;31m# Backpropagation\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[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\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[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\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[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\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;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/_tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    394\u001b[0m                 \u001b[0mcreate_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    395\u001b[0m                 inputs=inputs)\n\u001b[0;32m--> 396\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minputs\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    397\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    398\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\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/usr/local/lib/python3.8/dist-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    171\u001b[0m     \u001b[0;31m# some Python versions print out the first line of a multi-line function\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    172\u001b[0m     \u001b[0;31m# calls in the traceback and some print out the last line\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 173\u001b[0;31m     Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass\n\u001b[0m\u001b[1;32m    174\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    175\u001b[0m         allow_unreachable=True, accumulate_grad=True)  # Calls into the C++ engine to run the backward pass\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ],
      "source": [
        "encoder_model = encoder_model.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",
        "for t in range(100):\n",
        "    batch = next(iter(train_loader))\n",
        "    train(batch, encoder_model, loss_fn, optimizer)\n",
        "    if (t + 1) % 10 == 0:\n",
        "      test(X_train, y_train, edge_index_train, \"TRAIN\", encoder_model, loss_fn)\n",
        "test(X_test, y_test, edge_index_test, \"TEST\", encoder_model, loss_fn)\n",
        "print(\"Done!\")\n",
        "\n",
        "encoder = encoder_model[0]\n",
        "encoder.requires_grad=False\n",
        "\n",
        "# for name, param in encoder_model.named_parameters():\n",
        "#     if param.requires_grad:\n",
        "#         print(name, param.data)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## PMAT Training"
      ],
      "metadata": {
        "id": "WOoJSnrTSVWZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from pyvacy import optim, analysis\n",
        "# PMAT\n",
        "pmat_model = nn.Sequential(\n",
        "    encoder,\n",
        "    PMAT(K_hop, 60, agg_sigma),\n",
        "    Classification(K_hop, [], [(K_hop+1)*60, num_classes])\n",
        ")\n",
        "pmat_model = pmat_model.to(device)\n",
        "loss_fn = nn.CrossEntropyLoss()\n",
        "# optimizer = torch.optim.Adam(encoder_model.parameters(), lr=1e-3)\n",
        "optimizer = optim.DPAdam(\n",
        "    l2_norm_clip=1.0,\n",
        "    noise_multiplier=1.8,\n",
        "    batch_size=batch_size,\n",
        "    params=pmat_model.parameters(),\n",
        "    lr=0.5e-1)\n",
        "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)"
      ],
      "metadata": {
        "id": "LKxUUyl_SW-i"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for t in range(100):\n",
        "    batch = next(iter(train_loader))\n",
        "    train(batch, pmat_model, loss_fn, optimizer)\n",
        "    epsilon = analysis.moments_accountant(X_train.size(dim=0), batch_size, 1.8, t, pmat_delta)\n",
        "    if (t + 1) % 10 == 0:\n",
        "      test(X_train, y_train, edge_index_train, \"TRAIN\", pmat_model, loss_fn)\n",
        "      print(\"Optimizer Achieves ({:>0.1f}, {})-DP\".format(epsilon, pmat_delta))\n",
        "      print(\"LR:\", scheduler.get_last_lr()[0])\n",
        "    scheduler.step()\n",
        "    if epsilon >= pmat_epsilon:\n",
        "      break\n",
        "test(X_test, y_test, edge_index_test, \"TEST\", pmat_model, loss_fn)\n",
        "print(\"Done!\")\n",
        "\n",
        "pmat = pmat_model[1]\n",
        "pmat.requires_grad = False"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 294
        },
        "id": "CpsiubMeSoTE",
        "outputId": "cd9c8daf-e16f-4451-b595-c43a33c63ff2"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "error",
          "ename": "NameError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-1-1b1fec410ac2>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mepochs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m100\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\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[0;32m----> 3\u001b[0;31m     \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpmat_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      4\u001b[0m     \u001b[0mepsilon\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0manalysis\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmoments_accountant\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX_train\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mX_train\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1.8\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpmat_delta\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mt\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\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[0;31mNameError\u001b[0m: name 'train' is not defined"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zqf6_kLMUZUn"
      },
      "source": [
        "## Full Model Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2os_fvuNJmcJ"
      },
      "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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lo_SC_jUkl-y"
      },
      "source": [
        "Train full model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UPx__jwgLUCL"
      },
      "outputs": [],
      "source": [
        "model = GAP(encoder, \n",
        "            PMA(K_hop, agg_sigma), \n",
        "            Classification(K_hop, [60, 20], [(K_hop+1)*20, num_classes]))\n",
        "# model = GAP(encoder, \n",
        "#             PMAT(K_hop, 60, agg_sigma), \n",
        "#             Classification(K_hop, [60, 20], [(K_hop+1)*20, num_classes]))\n",
        "model = model.to(device)\n",
        "loss_fn = nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=1e-1)\n",
        "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.6)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z22WhxLcaveo",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "8a207048-a8e5-4cf9-f27f-8882c81fe072"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train Error: \n",
            " Accuracy: 50.0%, Loss: 1.086587\n",
            "Train Error: \n",
            " Accuracy: 70.8%, Loss: 0.908764\n",
            "Train Error: \n",
            " Accuracy: 81.5%, Loss: 0.759028\n",
            "Train Error: \n",
            " Accuracy: 85.6%, Loss: 0.639309\n",
            "Train Error: \n",
            " Accuracy: 88.4%, Loss: 0.542409\n",
            "Train Error: \n",
            " Accuracy: 89.9%, Loss: 0.495274\n",
            "Train Error: \n",
            " Accuracy: 91.3%, Loss: 0.453974\n",
            "Train Error: \n",
            " Accuracy: 91.9%, Loss: 0.418225\n",
            "Train Error: \n",
            " Accuracy: 92.4%, Loss: 0.387867\n",
            "Train Error: \n",
            " Accuracy: 92.7%, Loss: 0.362565\n",
            "Train Error: \n",
            " Accuracy: 93.0%, Loss: 0.350013\n",
            "Train Error: \n",
            " Accuracy: 93.1%, Loss: 0.337904\n",
            "Train Error: \n",
            " Accuracy: 93.4%, Loss: 0.326529\n",
            "Train Error: \n",
            " Accuracy: 93.5%, Loss: 0.316240\n",
            "Train Error: \n",
            " Accuracy: 93.5%, Loss: 0.306924\n",
            "Train Error: \n",
            " Accuracy: 93.5%, Loss: 0.301947\n",
            "Train Error: \n",
            " Accuracy: 93.7%, Loss: 0.297793\n",
            "Train Error: \n",
            " Accuracy: 93.8%, Loss: 0.293428\n",
            "Train Error: \n",
            " Accuracy: 93.7%, Loss: 0.288749\n",
            "Train Error: \n",
            " Accuracy: 93.8%, Loss: 0.284526\n",
            "Test Error: \n",
            " Accuracy: 90.0%, Loss: 0.391147\n",
            "Done!\n"
          ]
        }
      ],
      "source": [
        "for t in range(100):\n",
        "    batch = next(iter(train_loader))\n",
        "    train(batch, model, loss_fn, optimizer)\n",
        "    if (t + 1) % 10 == 0:\n",
        "      test(X_train, y_train, edge_index_train, \"TRAIN\", model, loss_fn) \n",
        "    scheduler.step()\n",
        "test(X_test, y_test, edge_index_test, \"TEST\", model, loss_fn)\n",
        "print(\"Done!\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 389
        },
        "id": "x0yD8RqskkhV",
        "outputId": "b9900a7b-b1ed-400a-de66-3d9a1ee3841e"
      },
      "outputs": [
        {
          "ename": "IndexError",
          "evalue": "ignored",
          "output_type": "error",
          "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": {
      "collapsed_sections": [
        "vb1hVZklFoDP",
        "GF264FWv0pJI",
        "Ty_dao27oEAB",
        "WebCohAOIX0P",
        "XDJP3FhOi_Ty",
        "E7oEXJ63oyl-",
        "BgX93DaHqqzh"
      ],
      "provenance": [],
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}