2-2-Stochastic training of GNN for link prediction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "rural-jenny",
   "metadata": {},
   "source": [
    "# Stochastic Training of GNN for Link Prediction\n",
    "This tutorial will show how to train a multi-layer GraphSAGE for link prediction on ogbn-arxiv provided by Open Graph Benchmark (OGB). The dataset contains around 170 thousand nodes and 1 million edges.\n",
    "\n",
    "By the end of this tutorial, you will be able to:\n",
    "\n",
    "    Train a GNN model for link prediction on a single GPU with DGL’s neighbor sampling components."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "hourly-cache",
   "metadata": {},
   "source": [
    "### Link Prediction Overview\n",
    "Link prediction requires the model to predict the probability of existence of an edge. This tutorial does so by computing a dot product between the representations of both incident nodes.\n",
    "\n",
    "\n",
    "It then minimizes the following binary cross entropy loss."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "excess-medicare",
   "metadata": {},
   "source": [
    "$\\hat{y}_{u\\sim v}=\\sigma(h_u^T h_v)$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "uniform-seventh",
   "metadata": {},
   "source": [
    "This is the formulation for generate loss for link prediction."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "backed-preference",
   "metadata": {},
   "source": [
    "$\\mathcal{L} = -\\sum_{u\\sim v\\in \\mathcal{D}}\\left( y_{u\\sim v}\\log(\\hat{y}_{u\\sim v}) + (1-y_{u\\sim v})\\log(1-\\hat{y}_{u\\sim v})) \\right)$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "thorough-translation",
   "metadata": {},
   "source": [
    "### Loading Dataset\n",
    "This tutorial loads the dataset from the ogb package.\n",
    "\n",
    "Ogbn-arxiv is a dataset for node classification. Here we use unsupervised link prediction training method to train a GNN model and generate node embeddings for the downstream task, i.e. the node classification task for ogbn-arxiv.\n",
    "\n",
    "Then we train a linear classifier for predicting the subject areas of arXiv CS papers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "occupied-italy",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dgl\n",
    "import torch\n",
    "import numpy as np\n",
    "from ogb.nodeproppred import DglNodePropPredDataset\n",
    "\n",
    "dataset = DglNodePropPredDataset('ogbn-arxiv')\n",
    "device = 'cpu'      # change to 'cuda' for GPU\n",
    "\n",
    "graph, node_labels = dataset[0]\n",
    "# Add reverse edges since ogbn-arxiv is unidirectional.\n",
    "graph = dgl.add_reverse_edges(graph)\n",
    "print(graph)\n",
    "print(node_labels)\n",
    "\n",
    "node_features = graph.ndata['feat']\n",
    "node_labels = node_labels[:, 0]\n",
    "num_features = node_features.shape[1]\n",
    "num_classes = (node_labels.max() + 1).item()\n",
    "print('Number of classes:', num_classes)\n",
    "\n",
    "idx_split = dataset.get_idx_split() # get default ogbn-arxiv data split\n",
    "train_nids = idx_split['train']\n",
    "valid_nids = idx_split['valid']\n",
    "test_nids = idx_split['test']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "excessive-paint",
   "metadata": {},
   "source": [
    "### Defining Neighbor Sampler and Data Loader in DGL\n",
    "Different from the link prediction for full graph, a common practice to train GNN on large graphs is to iterate over all the existing edges in minibatches, since computing the probability of all edges is usually impossible. For each minibatch of edges, you compute the output representation of their incident nodes using neighbor sampling and GNN, in a similar fashion introduced in the large-scale node classification tutorial.\n",
    "\n",
    "DGL provides ``dgl.dataloading.EdgeDataLoader`` to iterate over edges for edge classification or link prediction tasks.\n",
    "\n",
    "To perform link prediction, you need to specify a negative sampler. DGL provides builtin negative samplers such as ``dgl.dataloading.negative_sampler.Uniform``. Here this tutorial uniformly draws 5 negative examples per positive example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "prescription-reset",
   "metadata": {},
   "outputs": [],
   "source": [
    "negative_sampler = dgl.dataloading.negative_sampler.Uniform(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "married-tomato",
   "metadata": {},
   "source": [
    "After defining the negative sampler, one can then define the edge data loader with neighbor sampling. To create an ``EdgeDataLoader`` for link prediction, provide a neighbor sampler object as well as the negative sampler object created above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "champion-magazine",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = dgl.dataloading.MultiLayerNeighborSampler([4, 4])\n",
    "train_dataloader = dgl.dataloading.EdgeDataLoader(\n",
    "    # The following arguments are specific to NodeDataLoader.\n",
    "    graph,                                  # The graph\n",
    "    torch.arange(graph.number_of_edges()),  # The edges to iterate over\n",
    "    sampler,                                # The neighbor sampler\n",
    "    negative_sampler=negative_sampler,      # The negative sampler\n",
    "    device=device,                          # Put the MFGs on CPU or GPU\n",
    "    # The following arguments are inherited from PyTorch DataLoader.\n",
    "    batch_size=1024,    # Batch size\n",
    "    shuffle=True,       # Whether to shuffle the nodes for every epoch\n",
    "    drop_last=False,    # Whether to drop the last incomplete batch\n",
    "    num_workers=0       # Number of sampler processes\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "apparent-drink",
   "metadata": {},
   "source": [
    "You can peek one minibatch from ``train_dataloader`` and see what it will give you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pregnant-colon",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_nodes, pos_graph, neg_graph, mfgs = next(iter(train_dataloader))\n",
    "print('Number of input nodes:', len(input_nodes))\n",
    "print('Positive graph # nodes:', pos_graph.number_of_nodes(), '# edges:', pos_graph.number_of_edges())\n",
    "print('Negative graph # nodes:', neg_graph.number_of_nodes(), '# edges:', neg_graph.number_of_edges())\n",
    "print(mfgs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "toxic-supervision",
   "metadata": {},
   "source": [
    "The example minibatch consists of four elements.\n",
    "\n",
    "The first element is an ID tensor for the input nodes, i.e., nodes whose input features are needed on the first GNN layer for this minibatch.\n",
    "\n",
    "The second element is the positive graph for this minibatch\n",
    "\n",
    "The third element is the negative graph for this minibatch. In minibatch training, the positive graph and the negative graph only contain nodes necessary for computing the pair-wise scores of positive and negative examples in the current minibatch.\n",
    "\n",
    "The last element is a list of MFGs storing the computation dependencies for each GNN layer. The MFGs are used to compute the GNN outputs of the nodes involved in positive/negative graph."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "stopped-trade",
   "metadata": {},
   "source": [
    "### Defining Model for Node Representation\n",
    "The model is almost identical to the one in the previous node classification tutorial. The only difference is that since you are doing link prediction, the output dimension will not be the number of classes in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "quick-affect",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from dgl.nn import SAGEConv\n",
    "\n",
    "class Model(nn.Module):\n",
    "    def __init__(self, in_feats, h_feats):\n",
    "        super(Model, self).__init__()\n",
    "        self.conv1 = SAGEConv(in_feats, h_feats, aggregator_type='mean')\n",
    "        self.conv2 = SAGEConv(h_feats, h_feats, aggregator_type='mean')\n",
    "        self.h_feats = h_feats\n",
    "\n",
    "    def forward(self, mfgs, x):\n",
    "        h_dst = x[:mfgs[0].num_dst_nodes()]\n",
    "        h = self.conv1(mfgs[0], (x, h_dst))\n",
    "        h = F.relu(h)\n",
    "        h_dst = h[:mfgs[1].num_dst_nodes()]\n",
    "        h = self.conv2(mfgs[1], (h, h_dst))\n",
    "        return h\n",
    "\n",
    "model = Model(num_features, 128).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "helpful-cooperation",
   "metadata": {},
   "source": [
    "### Defining the Score Predictor for Edges\n",
    "After getting the node representation necessary for the minibatch, the last thing to do is to predict the score of the edges and non-existent edges in the sampled minibatch.\n",
    "\n",
    "The following score predictor is same as the one used in link prediction tutorial, takes a dot product between the incident nodes’ representations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "vital-commercial",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dgl.function as fn\n",
    "\n",
    "class DotPredictor(nn.Module):\n",
    "    def forward(self, g, h): # We have positive graph p_g and negative graph n_g\n",
    "        with g.local_scope():\n",
    "            g.ndata['h'] = h\n",
    "            # Compute a new edge feature named 'score' by a dot-product between the\n",
    "            # source node feature 'h' and destination node feature 'h'.\n",
    "            g.apply_edges(fn.u_dot_v('h', 'h', 'score'))\n",
    "            # u_dot_v returns a 1-element vector for each edge so you need to squeeze it.\n",
    "            return g.edata['score'][:, 0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "constitutional-custody",
   "metadata": {},
   "source": [
    "### Evaluating Performance\n",
    "There are various ways to evaluate the performance of link prediction. This tutorial follows the practice of GraphSAGE paper, where it treats the node embeddings learned by link prediction via training and evaluating a linear classifier on top of the learned node embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "employed-tyler",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generating node mebeddings of each node within the graph.\n",
    "def inference(model, graph, node_features):\n",
    "    with torch.no_grad():\n",
    "        nodes = torch.arange(graph.number_of_nodes())\n",
    "\n",
    "        sampler = dgl.dataloading.MultiLayerNeighborSampler([4, 4])\n",
    "        train_dataloader = dgl.dataloading.NodeDataLoader(\n",
    "            graph, torch.arange(graph.number_of_nodes()), sampler,\n",
    "            batch_size=1024,\n",
    "            shuffle=False,\n",
    "            drop_last=False,\n",
    "            num_workers=4,\n",
    "            device=device)\n",
    "\n",
    "        result = []\n",
    "        for input_nodes, output_nodes, mfgs in train_dataloader:\n",
    "            # feature copy from CPU to GPU takes place here\n",
    "            inputs = mfgs[0].srcdata['feat']\n",
    "            result.append(model(mfgs, inputs))\n",
    "\n",
    "        return torch.cat(result)\n",
    "\n",
    "import sklearn.metrics\n",
    "\n",
    "def evaluate(emb, label, train_nids, valid_nids, test_nids):\n",
    "    classifier = nn.Linear(emb.shape[1], num_classes).to(device) # linear classification layer\n",
    "    opt = torch.optim.LBFGS(classifier.parameters())\n",
    "\n",
    "    def compute_loss():\n",
    "        pred = classifier(emb[train_nids].to(device))\n",
    "        loss = F.cross_entropy(pred, label[train_nids].to(device))\n",
    "        return loss\n",
    "\n",
    "    def closure():\n",
    "        loss = compute_loss()\n",
    "        opt.zero_grad()\n",
    "        loss.backward()\n",
    "        return loss\n",
    "\n",
    "    prev_loss = float('inf')\n",
    "    for i in range(1000):\n",
    "        opt.step(closure)\n",
    "        with torch.no_grad():\n",
    "            loss = compute_loss().item()\n",
    "            if np.abs(loss - prev_loss) < 1e-4:\n",
    "                print('Converges at iteration', i)\n",
    "                break\n",
    "            else:\n",
    "                prev_loss = loss\n",
    "\n",
    "    with torch.no_grad():\n",
    "        pred = classifier(emb.to(device)).cpu()\n",
    "        label = label\n",
    "        valid_acc = sklearn.metrics.accuracy_score(label[valid_nids].numpy(), pred[valid_nids].numpy().argmax(1))\n",
    "        test_acc = sklearn.metrics.accuracy_score(label[test_nids].numpy(), pred[test_nids].numpy().argmax(1))\n",
    "    return valid_acc, test_acc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cutting-christopher",
   "metadata": {},
   "source": [
    "### Defining Training Loop\n",
    "The following initializes the model and defines the optimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "standard-academy",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model(node_features.shape[1], 128).to(device)\n",
    "predictor = DotPredictor().to(device)\n",
    "opt = torch.optim.Adam(list(model.parameters()) + list(predictor.parameters()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "coated-series",
   "metadata": {},
   "source": [
    "The following is the training loop for link prediction and evaluation, and also saves the model that performs the best on the validation set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "intense-order",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm\n",
    "import sklearn.metrics\n",
    "\n",
    "best_accuracy = 0\n",
    "best_model_path = 'model.pt'\n",
    "for epoch in range(1):\n",
    "    with tqdm.tqdm(train_dataloader) as tq:\n",
    "        for step, (input_nodes, pos_graph, neg_graph, mfgs) in enumerate(tq):\n",
    "            # feature copy from CPU to GPU takes place here\n",
    "            inputs = mfgs[0].srcdata['feat']\n",
    "\n",
    "            outputs = model(mfgs, inputs)\n",
    "            pos_score = predictor(pos_graph, outputs)\n",
    "            neg_score = predictor(neg_graph, outputs)\n",
    "\n",
    "            score = torch.cat([pos_score, neg_score])\n",
    "            label = torch.cat([torch.ones_like(pos_score), torch.zeros_like(neg_score)])\n",
    "            loss = F.binary_cross_entropy_with_logits(score, label)\n",
    "\n",
    "            opt.zero_grad()\n",
    "            loss.backward()\n",
    "            opt.step()\n",
    "\n",
    "            tq.set_postfix({'loss': '%.03f' % loss.item()}, refresh=False)\n",
    "\n",
    "            if (step + 1) % 500 == 0:\n",
    "                model.eval()\n",
    "                emb = inference(model, graph, node_features)\n",
    "                valid_acc, test_acc = evaluate(emb, node_labels, train_nids, valid_nids, test_nids)\n",
    "                print('Epoch {} Validation Accuracy {} Test Accuracy {}'.format(epoch, valid_acc, test_acc))\n",
    "                if best_accuracy < valid_acc:\n",
    "                    best_accuracy = valid_acc\n",
    "                    torch.save(model.state_dict(), best_model_path)\n",
    "                model.train()\n",
    "\n",
    "                # Note that this tutorial do not train the whole model to the end.\n",
    "                break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "prospective-teacher",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_latest_p36",
   "language": "python",
   "name": "conda_pytorch_latest_p36"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}