assignment01_contrastive_learning/assignment.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XBQ93q_yNfW7"
   },
   "source": [
    "# Assignment — Graph Contrastive Learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GeilxQ7QQGLq"
   },
   "source": [
    "### Task 1. Augmentation (2 points)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4_XqkYMlw0Nh"
   },
   "outputs": [],
   "source": [
    "#!pip install dgl-cu111 -f https://data.dgl.ai/wheels/repo.html -q\n",
    "!pip install dgl -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "XlYO66rKMLBR",
    "outputId": "1c4106d7-4042-4ef6-9a66-0fb97816f983"
   },
   "outputs": [],
   "source": [
    "import networkx as nx\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import requests\n",
    "\n",
    "import dgl\n",
    "from dgl.nn import GraphConv\n",
    "from dgl.dataloading import GraphDataLoader\n",
    "from dgl.data import DGLDataset\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.optim import Adam\n",
    "from torch.utils.data import Subset\n",
    "\n",
    "from sklearn.manifold import TSNE\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "\n",
    "from IPython.display import clear_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7b1jiI9ge3vs"
   },
   "source": [
    "<img src='https://raw.githubusercontent.com/netspractice/advanced_gnn/made2021/assignment_contrastive_learning/contrastive_learning.png' width=500>\n",
    "\n",
    "Source: https://arxiv.org/abs/2103.00111"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C0m1ts4NdaKL"
   },
   "source": [
    "Contrastive learning aims to learn representations by maximizing feature consistency under differently augmented views, that exploit data- or task-specific augmentations. In a case of graph representation learning, there are some augmentation techniques that can be used to produce graph embeddings for downstream tasks, say classification.\n",
    "\n",
    "Write a class `GraphAugmentation` with a function `transform` that takes a graph and returns an augmented graph. Types of augmentation:\n",
    "* `drop_nodes` — randomly drops a share of nodes with a given ratio\n",
    "* `pert_edges` — randomly perturbs (rewires) a share of edges with a given ratio\n",
    "* `attr_mask` — randomly masks a share of node attributes with a given ratio and a name in `ndata` collection\n",
    "* `rw_subgraph` — builds a subgraph based on random walk\n",
    "* `identical` — the same graph, no augmentation\n",
    "\n",
    "Augmentations are applied to graphs with self-loops, so keep self-loops during edges perturbation. Parallel edges are allowed after perturbation. A random walk subgraph is constructed by (1) adding a random starting node, (2) adding all its neighbors, (3) adding all neighbors of a random node in the subgraph and repeating the step 3 while number of nodes exceeds the threshold `(1 - ratio)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "id": "NPa_dB39Hc-8",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "2c784c04ab9133ba352688dc0d28f157",
     "grade": false,
     "grade_id": "cell-4aec214abdad2607",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "class GraphAugmentation():\n",
    "    def __init__(self, type, ratio=0.2, node_feat='attr'):\n",
    "        self.type = type\n",
    "        self.ratio = ratio\n",
    "        self.node_feat = node_feat\n",
    "    \n",
    "    def transform(self, g):\n",
    "        if self.type == 'drop_nodes':\n",
    "            return self.drop_nodes(g)\n",
    "        elif self.type == 'pert_edges':\n",
    "            return self.pert_edges(g)\n",
    "        elif self.type == 'attr_mask':\n",
    "            return self.attr_mask(g)\n",
    "        elif self.type == 'rw_subgraph':\n",
    "            return self.rw_subgraph(g)\n",
    "        elif self.type == 'identical':\n",
    "            return g\n",
    "    \n",
    "    def drop_nodes(self, g):\n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def pert_edges(self, g):\n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def attr_mask(self, g):\n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def rw_subgraph(self, g):\n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "CThFcpnILmQ8",
    "outputId": "ddca390c-5f91-47c4-f951-2022997c02c9"
   },
   "outputs": [],
   "source": [
    "g = dgl.rand_graph(100, 300)\n",
    "g = g.remove_self_loop()\n",
    "g = g.add_self_loop()\n",
    "g.ndata['attr'] = torch.ones(100, 10)\n",
    "g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "id": "7Ij-mR7QnTAg",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "9e8645e7198498d259e47d25d12bc4f1",
     "grade": true,
     "grade_id": "cell-54be8c4ec2511ae1",
     "locked": true,
     "points": 0.5,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "ratio = 0.7\n",
    "aug = GraphAugmentation('drop_nodes', ratio=ratio)\n",
    "aug_g = aug.transform(g)\n",
    "assert aug_g.ndata['attr'].shape[1] == g.ndata['attr'].shape[1]\n",
    "assert aug_g.ndata['attr'].shape[0] < g.ndata['attr'].shape[0]\n",
    "assert aug_g.ndata['attr'].shape[0] == int(g.number_of_nodes() * (1 - ratio))\n",
    "G = nx.Graph(aug_g.to_networkx())\n",
    "assert np.isclose(nx.laplacian_spectrum(G), 0).sum() > 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "id": "l_AEnNPQiRxJ",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "b7764edf0b7d65b11a23519f6889b420",
     "grade": true,
     "grade_id": "cell-9ac3eba85f9b9ab9",
     "locked": true,
     "points": 0.5,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "aug = GraphAugmentation('pert_edges', ratio=0.2)\n",
    "aug_g = aug.transform(g)\n",
    "assert aug_g.ndata['attr'].shape == g.ndata['attr'].shape\n",
    "assert aug_g.number_of_edges() == g.number_of_edges()\n",
    "assert not torch.all(aug_g.adj().to_dense() == g.adj().to_dense())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "id": "f3ITr6Grp9aJ",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "d129a902df268c14086a67e812365d05",
     "grade": true,
     "grade_id": "cell-7d532cfaeeac9a18",
     "locked": true,
     "points": 0.5,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "aug = GraphAugmentation('attr_mask', ratio=0.2, node_feat='attr')\n",
    "aug_g = aug.transform(g)\n",
    "assert aug_g.ndata['attr'].shape == (100, 10)\n",
    "mask = (aug_g.ndata['attr'][0, :] == 0).repeat(100, 1)\n",
    "assert torch.all(aug_g.ndata['attr'][mask] == 0)\n",
    "assert torch.all(aug_g.ndata['attr'][~mask] == 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "id": "hpeaqyngyOml",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "28b3920ebb28069a2df4d8dd5ee6fbf1",
     "grade": true,
     "grade_id": "cell-255aea5b0cba7534",
     "locked": true,
     "points": 0.5,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "aug = GraphAugmentation('rw_subgraph', ratio=0.7)\n",
    "aug_g = aug.transform(g)\n",
    "assert aug_g.ndata['attr'].shape[1] == g.ndata['attr'].shape[1]\n",
    "assert aug_g.ndata['attr'].shape[0] < g.ndata['attr'].shape[0]\n",
    "G = nx.Graph(aug_g.to_networkx())\n",
    "assert np.isclose(nx.laplacian_spectrum(G), 0).sum() == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pwRRJ4cHQBN5"
   },
   "source": [
    "### Task 2. Contrastive dataset (2 points)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C6K-_4ny_Mj-"
   },
   "source": [
    "We will fed augmented graphs into encoder during training to obtain graph embeddings. Let us prepare a graph contrastive dataset class so that each element in the dataset will represent augmented graphs and a label.\n",
    "\n",
    "Write a class `ContrastiveDataset` with a function `__getitem__` that takes a graph's index and returns a tuple:\n",
    "* an initial graph\n",
    "* a graph after the first augmentation\n",
    "* a graph after the second augmentation\n",
    "* a label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "id": "as-r-yeMCJNo",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "0b35f0d1b411c9a660dcbf52badbc5a1",
     "grade": false,
     "grade_id": "cell-8467bac5d30fb623",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "class ContrastiveDataset(DGLDataset):\n",
    "    def __init__(self, filename, augmentations):\n",
    "        self.filename = filename\n",
    "        self.graphs = None\n",
    "        self.labels = None\n",
    "        self.augmentations = augmentations\n",
    "        assert len(self.augmentations) == 2\n",
    "        super().__init__(name=filename)\n",
    "\n",
    "    def process(self):\n",
    "        graphs, graph_data = dgl.load_graphs(self.filename)\n",
    "        self.graphs = graphs\n",
    "        self.labels = graph_data['labels']\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.graphs)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uar5OeYrXrUb"
   },
   "source": [
    "PROTEINS is a dataset with 1113 proteins where nodes are secondary structure elements and there is an edge between two nodes if they are neighbors in the amino-acid sequence or in 3D space. It has 3 discrete labels, representing helix, sheet or turn. Proteins are divided into two classes: enzymes and non-enzymes. Source: https://arxiv.org/abs/2007.08663.\n",
    "\n",
    "Let us create a dataset with dropping nodes and masking attributes augmentations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0F8hCtq3QLGr",
    "outputId": "73326f42-37d2-4c60-da33-51c478049fc0"
   },
   "outputs": [],
   "source": [
    "url = 'https://github.com/netspractice/advanced_gnn/raw/made2021/assignment_contrastive_learning/proteins.bin'\n",
    "open('proteins.bin', 'wb').write(requests.get(url).content);\n",
    "\n",
    "augmentations = []\n",
    "augmentations.append(GraphAugmentation('drop_nodes', ratio=0.1))\n",
    "augmentations.append(GraphAugmentation('attr_mask', ratio=0.4, node_feat='attr'))\n",
    "dataset = ContrastiveDataset(filename='proteins.bin', augmentations=augmentations)\n",
    "N = len(dataset)\n",
    "N"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "id": "llq6kj49TDwk",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "934b70d94108f159148afa825fc897a1",
     "grade": true,
     "grade_id": "cell-2cdb6b25f5ca277f",
     "locked": true,
     "points": 1,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "g, aug_g1, aug_g2, label = dataset[0]\n",
    "assert g.ndata['attr'].shape == (42, 3)\n",
    "assert aug_g1.ndata['attr'].shape == (37, 3)\n",
    "assert g.ndata['attr'].sum() <= 42"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qRa9FMz7Ui2G"
   },
   "source": [
    "Since we perform the random walk subgraph augmentation, we want to make sure all initial graphs are connected.\n",
    "\n",
    "Write a function `connected_subset` that takes an initial dataset and returns a `torch.utils.data.dataset.Subset` with connected graphs only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "id": "pIV9YIEBgBXq",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "52f536a8f39dec265afee6a74e60841b",
     "grade": false,
     "grade_id": "cell-5a6861794b797604",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def connected_subset(dataset):\n",
    "    # YOUR CODE HERE\n",
    "    raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "id": "D2VoZg8covCb",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "9fa5169511bc79aeeac4908743c1551f",
     "grade": true,
     "grade_id": "cell-b6e99b0412f848fe",
     "locked": true,
     "points": 1,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "c_dataset = connected_subset(dataset)\n",
    "N = len(c_dataset)\n",
    "assert N == 1067"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "P7EnVHvEgtjR"
   },
   "source": [
    "Let us look at some graphs in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 699
    },
    "id": "XYZnTQPhldOo",
    "outputId": "c8c7c71e-82de-4079-df8d-93e0755529a9"
   },
   "outputs": [],
   "source": [
    "colors = ['tab:orange', 'tab:green']\n",
    "plt.figure(figsize=(12, 12))\n",
    "np.random.seed(0)\n",
    "for i in range(16):\n",
    "    plt.subplot(4, 4, i+1)\n",
    "    g, _, _, l = c_dataset[np.random.randint(N)]\n",
    "    g = nx.Graph(g.to_networkx())\n",
    "    g.remove_edges_from(nx.selfloop_edges(g))\n",
    "    nx.draw_kamada_kawai(g, node_size=30, node_color=colors[l])\n",
    "    plt.title('enzymes' if l == 1 else 'non-enzymes')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "84QjNrL8i6SZ"
   },
   "source": [
    "### Task 3. GCN Encoder (1 point)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FI_iKY-_jCQS"
   },
   "source": [
    "Let an encoder be the two-layers GCN (`GraphConv` in `dgl`) with mean graph pooling and two-layers MLP projection head. All layers except of input and output ones have `hidden_dim` dimensionality. Apply ReLU as an activation function.\n",
    "\n",
    "Write a class `GCNEncoder` with a function `forward` that takes a batch of graphs, node attrubute name in `ndata` collection and returns graph embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "id": "6pZhje4sVARU",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "2554688040d350a2271e07f5de9e6b32",
     "grade": false,
     "grade_id": "cell-ee3f82340b311a08",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "class GCNEncoder(nn.Module):\n",
    "    def __init__(self, input_dim, hidden_dim, output_dim):\n",
    "        super().__init__()\n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def forward(self, g, node_feat):\n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "id": "SCB99OXNlqyt",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "57cb0d13204c52aa381970160d14d798",
     "grade": true,
     "grade_id": "cell-db01468c075a48bc",
     "locked": true,
     "points": 1,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "batch = []\n",
    "for _ in range(2):\n",
    "    g = dgl.rand_graph(100, 300)\n",
    "    g = g.remove_self_loop()\n",
    "    g = g.add_self_loop()\n",
    "    g.ndata['attr'] = torch.ones(100, 10)\n",
    "    batch.append(g)\n",
    "batch = dgl.batch(batch)\n",
    "\n",
    "encoder = GCNEncoder(input_dim=10, hidden_dim=32, output_dim=16)\n",
    "emb = encoder(batch, 'attr')\n",
    "assert emb.shape == (2, 16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9iYuvugOno0F"
   },
   "source": [
    "### Task 4. Classification on untrained encoder (1 point)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "CN-HrVjDuT8a",
    "outputId": "433ec692-61ae-44b0-a032-6ab4207274b5"
   },
   "outputs": [],
   "source": [
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bstkkD0Dq8x1"
   },
   "source": [
    "Let us check the logistic regression model on the untrained encoder output.\n",
    "\n",
    "Write a function `train_test_split` that splits the dataset into train and test sets by given one-hot encoded vectors `train_idx` and `test_idx`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "id": "7XfwhKQ31Xb0",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "7df516a3472522bfebf761a874acc26c",
     "grade": false,
     "grade_id": "cell-8cf67c5878b238d1",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def train_test_split(c_dataset, train_idx, test_idx):\n",
    "    # YOUR CODE HERE\n",
    "    raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "id": "bsDjcmsHrzFg",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "75d9bbc1938183c02af3d2c7ba89d2df",
     "grade": true,
     "grade_id": "cell-a502a064c530464e",
     "locked": true,
     "points": 1,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(0)\n",
    "ratio = [0.9, 0.1] # train test ratio\n",
    "split_idx = ['train'] * int(ratio[0] * N) \\\n",
    "    + ['test'] * int(ratio[1] * N)\n",
    "split_idx = np.random.permutation(split_idx)\n",
    "train_idx = np.where(split_idx == 'train')[0]\n",
    "test_idx = np.where(split_idx == 'test')[0]\n",
    "\n",
    "graph_train, graph_test, y_train, y_test = train_test_split(\n",
    "    c_dataset, train_idx, test_idx)\n",
    "assert graph_train.ndata['attr'].shape == (36363, 3)\n",
    "assert graph_test.ndata['attr'].shape == (3740, 3)\n",
    "assert y_train.shape == (960, )\n",
    "assert y_test.shape == (106, )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gucFvSQ_tmAn"
   },
   "source": [
    "Let us check the classification score and look at tSNE visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "auBC65L3EOqj",
    "outputId": "83ff94c9-72e1-4263-e444-da15f62ce84d"
   },
   "outputs": [],
   "source": [
    "encoder = GCNEncoder(input_dim=3, hidden_dim=32, output_dim=16)\n",
    "encoder.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "szSDScXGtKWn"
   },
   "outputs": [],
   "source": [
    "def classification_score(graph_train, graph_test, y_train, y_test, encoder, show=True):\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        X_train = encoder(graph_train, 'attr').cpu()\n",
    "        X_test = encoder(graph_test, 'attr').cpu()\n",
    "    \n",
    "    clf = LogisticRegression()\n",
    "    clf.fit(X_train, y_train.cpu())\n",
    "    score = clf.score(X_test, y_test.cpu())\n",
    "    \n",
    "    if show:\n",
    "        plt.figure(figsize=(10, 6))\n",
    "        cmap = plt.cm.Set1_r\n",
    "        dec = TSNE(n_components=2)\n",
    "        xy_emb = dec.fit_transform(X_train)\n",
    "        plt.scatter(xy_emb[:, 0], xy_emb[:, 1], c=y_train.cpu(), cmap=cmap, s=5)\n",
    "        plt.title('tSNE visualization')\n",
    "        plt.show()\n",
    "        print('Accuracy: {:.4f}'.format(score))\n",
    "    \n",
    "    return score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 407
    },
    "id": "ui3w4WhUoILR",
    "outputId": "2e98bd00-6a03-468b-d63a-da5d6f4cb9e7"
   },
   "outputs": [],
   "source": [
    "score = classification_score(graph_train, graph_test, y_train, y_test, encoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xQTjLzXCvbY7"
   },
   "source": [
    "### Task 5. Contrastive loss (2 points)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "N-wiMTo1vidM"
   },
   "source": [
    "In graph contrastive learning, pre-training is performed through maximizing the agreement between two augmented views of the same graph via a contrastive loss in the latent space. Contrastive loss function is defined to enforce maximizing the consistency between positive pairs $z_i$, $z_j$ (the same graph under different augmentations) compared with negative pairs. Here we utilize the NT-Xent Loss that is defined for $n$-th graph in a batch of $N$ graphs as follows:\n",
    "\n",
    "$$l_{n}=-\\log \\frac{\\exp \\left(\\text{sim}(z_{n, i}, z_{n, j}) / \\tau \\right)}{\\sum_{n'=1, n' \\neq n}^{N} \\exp \\left( \\text{sim}(z_{n, i}, z_{n', j}) / \\tau \\right)}$$\n",
    "\n",
    "where $\\text{sim}$ is cosine similarity $\\text{sim}(z_i, z_j) = z_i^\\top z_j / (\\| z_i \\| \\cdot \\| z_j \\|)$ and $\\tau$ is a temperature parameter.\n",
    "\n",
    "Source: https://arxiv.org/pdf/2010.13902.pdf\n",
    "\n",
    "Write a function `ntxent` that takes a batch of agmented graph embeddings `x1` and a batch of agmented graph embeddings `x2` and returns mean loss value among all graphs $L = \\frac{1}{N}\\sum_{n=1}^N l_n$.\n",
    "\n",
    "_Hint: it is possible to use matrix operations only, with no loops._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "id": "8iVXKw0FPRxF",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "4a33dd874d1f5b4dccb061d38ab9f570",
     "grade": false,
     "grade_id": "cell-0a6e245edc54a475",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def ntxent(x1, x2, tau=0.1):\n",
    "    # YOUR CODE HERE\n",
    "    raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "id": "hTvZzSYq0hri",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "baa94fdb8c3f72676e661d6a001abab6",
     "grade": true,
     "grade_id": "cell-cb7ca97cf5793008",
     "locked": true,
     "points": 2,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "x1 = torch.tensor([[1., 0.], [0., 1.]])\n",
    "x2 = torch.tensor([[1., 0.], [0., 1.]])\n",
    "assert ntxent(x1, x2) == -10\n",
    "\n",
    "x1 = torch.tensor([[1., 0.], [0., 1.]])\n",
    "x2 = torch.tensor([[0., 1.], [1., 0.]])\n",
    "assert ntxent(x1, x2) == 10\n",
    "\n",
    "torch.manual_seed(0)\n",
    "x1 = torch.randn(128, 16)\n",
    "x2 = torch.randn(128, 16)\n",
    "assert round(ntxent(x1, x2).item(), 4) == 7.191"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oo1nDTaw2xcY"
   },
   "source": [
    "### Task 6. Training loop (1 point)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_i2kLGBb3R1y"
   },
   "source": [
    "Let us train encoder under contrastive loss and then check classification score."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "N_7c65uya8th"
   },
   "outputs": [],
   "source": [
    "loader = GraphDataLoader(\n",
    "    c_dataset,\n",
    "    batch_size=64,\n",
    "    drop_last=False,\n",
    "    shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3YUqmdm6oMtZ"
   },
   "outputs": [],
   "source": [
    "encoder = GCNEncoder(input_dim=3, hidden_dim=32, output_dim=16)\n",
    "encoder.to(device)\n",
    "opt = Adam(encoder.parameters(), lr=0.005)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "59m12kJ44Phx"
   },
   "source": [
    "Write a function `train` that takes augmented batches, makes optimization step and returns a loss value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "id": "gesyjPE03qYB",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "e1d681468270b7a9fdcfb1a46ec5ff5d",
     "grade": false,
     "grade_id": "cell-927cbd6e46929578",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def train(encoder, aug_batch1, aug_batch2, opt):\n",
    "    # YOUR CODE HERE\n",
    "    raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "id": "WdL1hIKg445o",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "5443cd13e0cea4238e275374b4054da6",
     "grade": true,
     "grade_id": "cell-0f733c004d847d5d",
     "locked": true,
     "points": 0,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "for batch, aug_batch1, aug_batch2, label in loader:\n",
    "    break\n",
    "loss_item = train(encoder, aug_batch1, aug_batch2, opt)\n",
    "assert type(loss_item) == float\n",
    "assert loss_item > 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2D7g--Sp6xuu"
   },
   "source": [
    "Here is a training loop that accumulates mean loss per epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 281
    },
    "id": "YxCp6MPtEjjK",
    "outputId": "731abffb-d6db-41d4-e606-6a4e958e79b4"
   },
   "outputs": [],
   "source": [
    "loss_vals = []\n",
    "n_epochs = 30\n",
    "for i in range(n_epochs):\n",
    "    loss_epoch = []\n",
    "    for batch, aug_batch1, aug_batch2, label in loader:\n",
    "        loss_item = train(encoder, aug_batch1, aug_batch2, opt)\n",
    "        loss_epoch.append(loss_item)\n",
    "    loss_vals.append(sum(loss_epoch)/len(loss_epoch))\n",
    "    plt.plot(loss_vals)\n",
    "    plt.title('Contrastive loss. Epoch: {}/{}'.format(i+1, n_epochs))\n",
    "    plt.show();\n",
    "    clear_output(wait=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 407
    },
    "deletable": false,
    "editable": false,
    "id": "0jCVRYCVhLcz",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "49e14ce5bf70181ba3f14e3c78ea9524",
     "grade": true,
     "grade_id": "cell-fb08ba989918e2e7",
     "locked": true,
     "points": 1,
     "schema_version": 3,
     "solution": false,
     "task": false
    },
    "outputId": "e8c45489-5dae-4ef5-b075-24131989af46"
   },
   "outputs": [],
   "source": [
    "score = classification_score(graph_train, graph_test, y_train, y_test, encoder)\n",
    "assert score > 0.65"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z-IdnnaR-9Yl"
   },
   "source": [
    "As we see, we can noticeably improve classification score using self-supervised learning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kOUZh6Ju_Ng0"
   },
   "source": [
    "### Task 7. Augmentation comparison (1 point)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3pVfWywCAEOe"
   },
   "source": [
    "Here we aim to compare augmentation techniques and conclude which pair is better for PROTEINS dataset.\n",
    "\n",
    "Write a function `run` that takes a filename with proteins, number of epochs and a list of types of augmentation. It returns a np.array with a classification score matrix where rows are first augmentation, columns are second augmentation. Since the matrix is asymptotically symmetric, calculate the upper triangle values only.\n",
    "\n",
    "It can take time. To speed up the evaluation, return calculated score matrix without actual training:\n",
    "```\n",
    "def run(filename, n_epochs, augs):\n",
    "    scores = [[0.5, 0.5, 0.5, 0.5, 0.5], [0, 0.5, 0.5, 0.5], ...\n",
    "    return scores\n",
    "\n",
    "    ### ACTUAL TRAINING\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "id": "2zz-NLC4ADDb",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "337480ceafaf757f7e6fe81f86e4985b",
     "grade": false,
     "grade_id": "cell-f63d23d10982e493",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def run(filename, n_epochs, augs):\n",
    "    # YOUR CODE HERE\n",
    "    raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "id": "aZKXihALD-bC",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "9bacd4e482c685212ad3ce17189341fe",
     "grade": true,
     "grade_id": "cell-bbe64c2a5aac60d2",
     "locked": true,
     "points": 1,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "augs = ['drop_nodes', 'pert_edges', 'attr_mask', 'rw_subgraph', 'identical']\n",
    "res = run('proteins.bin', n_epochs=30, augs=augs)\n",
    "symm = (res.T + res - np.diag(res[range(5), range(5)])).sum(0)\n",
    "assert np.all((res > 0).sum(0) == np.arange(5) + 1)\n",
    "assert np.all(res[res > 0] > 0.6)\n",
    "assert symm[0] > symm[4]\n",
    "assert symm[2] > symm[4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 204
    },
    "id": "hfBiE7e7ElW8",
    "outputId": "f92ff9a4-8cd3-4d4b-9d83-c6eaab7ca486"
   },
   "outputs": [],
   "source": [
    "pd.DataFrame(res, index=augs, columns=augs).round(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5uQrdLjaWZ5Z"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "assignment.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}