{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright (c) 2022 Graphcore Ltd. All rights reserved."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Temporal Graph Networks training\n",
    "This notebook demonstrates how to train [Temporal Graph Networks](https://arxiv.org/abs/2006.10637)\n",
    "(TGNs) on the IPU. See our [blog post](https://www.graphcore.ai/posts/accelerating-and-scaling-temporal-graph-networks-on-the-graphcore-ipu) for details on\n",
    "performance and scaling.\n",
    "TGN can be used to predict connections in a dynamically evolving graph. This application looks at graphs that gain edges over time. A typical use case is a social network where users form new connections over time, or a recommendation system where new edges represent an interaction of a user with a product or content.\n",
    "\n",
    "![dynamic_graph.png](static/dynamic_graph.png)\n",
    "\n",
    "In this notebook we apply TGN to the [JODIE Wikipedia dataset](https://snap.stanford.edu/jodie/), a dynamic graph of 1,000 Wikipedia articles and 8,227  Wikipedia users. 157,474 time-stamped edges describe the interactions of users with articles."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running on Paperspace\n",
    "\n",
    "The Paperspace environment lets you run this notebook with no set up. To improve your experience we preload datasets and pre-install packages, this can take a few minutes, if you experience errors immediately after starting a session please try restarting the kernel before contacting support. If a problem persists or you want to give us feedback on the content of this notebook, please reach out to through our community of developers using our [slack channel](https://www.graphcore.ai/join-community) or raise a [GitHub issue](https://github.com/gradient-ai/Graphcore-Pytorch/issues)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "%pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb3322bc",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "from examples_utils import notebook_logging\n",
    "%load_ext gc_logger"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load required dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import poptorch\n",
    "import time\n",
    "import numpy as np\n",
    "from poptorch import DataLoader\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython import display\n",
    "from sklearn.metrics import average_precision_score, roc_auc_score\n",
    "\n",
    "from tgn_modules import TGN, Data, init_weights, DataWrapper"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the model\n",
    "In addition to a single layer attention-based graph neural network (`tgn_gnn`) TGN\n",
    "introduces a memory module (`tgn_memory`) that keeps track of past interaction of each\n",
    "node.\n",
    "\n",
    "![architecture.png](static/architecture.png)\n",
    "\n",
    "Due to the dynamic nature of the graph, lower batch sizes will yield a higher accuracy.\n",
    "The hyperparameters `nodes_size` and `edges_size` control the padding of the tensors\n",
    "that contain the relevant nodes/edges per batch. The optimal setting for these\n",
    "hyperparameters depend on the batch size."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "LEARNING_RATE = 0.000075\n",
    "BATCH_SIZE = 40\n",
    "NODES_SIZE = 400\n",
    "EDGES_SIZE = 1200\n",
    "DEVICE_ITERATIONS = 212\n",
    "EPOCHS = 25"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create and initialise the TGN model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "tgn = TGN(\n",
    "    num_nodes=9227,\n",
    "    raw_msg_dim=172,\n",
    "    memory_dim=100,\n",
    "    time_dim=100,\n",
    "    embedding_dim=100,\n",
    "    dtype=torch.float32,\n",
    "    dropout=0.1,\n",
    "    target=\"ipu\",\n",
    ")\n",
    "tgn.apply(init_weights);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Dataloader\n",
    "We create an [IterableDataset](https://pytorch.org/docs/stable/data.html#torch.utils.data.IterableDataset) yielding batches of nodes, edges and negative samples, padded to a constant size and pass this to the [PopTorch DataLoader](https://docs.graphcore.ai/projects/poptorch-user-guide/en/latest/batching.html#poptorch-dataloader). By using `DEVICE_ITERATIONS > 1` and offloading the data loading process to a separate thread with `poptorch.DataLoaderMode.Async` we reduce the host overhead for data preprocessing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = DataWrapper(\n",
    "    Data(\"data/JODIE\", torch.float32, BATCH_SIZE, NODES_SIZE, EDGES_SIZE), \"train\"\n",
    ")\n",
    "test_data = DataWrapper(\n",
    "    Data(\"data/JODIE\", torch.float32, BATCH_SIZE, NODES_SIZE, EDGES_SIZE), \"val\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_opts = poptorch.Options()\n",
    "train_opts.deviceIterations(DEVICE_ITERATIONS)\n",
    "test_opts = poptorch.Options()\n",
    "test_opts.deviceIterations(1)\n",
    "\n",
    "torch.multiprocessing.set_sharing_strategy(\"file_system\")\n",
    "\n",
    "async_options = {\n",
    "    \"sharing_strategy\": poptorch.SharingStrategy.SharedMemory,\n",
    "    \"load_indefinitely\": True,\n",
    "    \"early_preload\": True,\n",
    "    \"buffer_size\": 2,\n",
    "}\n",
    "\n",
    "train_dl = DataLoader(\n",
    "    options=train_opts,\n",
    "    dataset=train_data,\n",
    "    batch_size=1,\n",
    "    mode=poptorch.DataLoaderMode.Async,\n",
    "    async_options=async_options,\n",
    ")\n",
    "test_dl = DataLoader(options=test_opts, dataset=test_data, batch_size=1)\n",
    "\n",
    "dataset_size = len(train_dl) * (DEVICE_ITERATIONS * BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare the model\n",
    "Define the [PopTorch optimizer](https://docs.graphcore.ai/projects/poptorch-user-guide/en/latest/pytorch_to_poptorch.html#optimizers) and wrap the PyTorch model into a [PopTorch tranining and inference model](https://docs.graphcore.ai/projects/poptorch-user-guide/en/latest/pytorch_to_poptorch.html#creating-your-model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "optim = poptorch.optim.AdamW(\n",
    "    tgn.parameters(),\n",
    "    lr=LEARNING_RATE,\n",
    "    bias_correction=True,\n",
    "    weight_decay=0.0,\n",
    "    eps=1e-8,\n",
    "    betas=(0.9, 0.999),\n",
    ")\n",
    "\n",
    "tgn_train = poptorch.trainingModel(tgn, options=train_opts, optimizer=optim)\n",
    "tgn_eval = poptorch.inferenceModel(tgn, options=test_opts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define a training and inference loop and train TGN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_train(model, train_data, do_reset) -> float:\n",
    "    \"\"\"Trains TGN for one epoch\"\"\"\n",
    "    total_loss = 0\n",
    "    num_events = 0\n",
    "    model.train()\n",
    "\n",
    "    if do_reset:\n",
    "        model.memory.reset_state()  # Start with a fresh memory.\n",
    "        model.copyWeightsToDevice()\n",
    "    for n, batch in enumerate(train_data):\n",
    "        count, loss = model(**batch)\n",
    "        total_loss += float(loss) * count\n",
    "        num_events += count\n",
    "        model.memory.detach()\n",
    "\n",
    "    return total_loss / num_events\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def run_test(model, inference_data) -> (float, float):\n",
    "    \"\"\"Inference over one epoch\"\"\"\n",
    "\n",
    "    model.eval()\n",
    "    torch.manual_seed(12345)  # Ensure deterministic sampling across epochs\n",
    "    aps = 0.0\n",
    "    aucs = 0.0\n",
    "    num_events = 0\n",
    "    for batch in inference_data:\n",
    "        count, y_true, y_pred = model(**batch)\n",
    "        aps += count * average_precision_score(y_true, y_pred)\n",
    "        aucs += count * roc_auc_score(y_true, y_pred)\n",
    "        num_events += count\n",
    "    return aps / num_events, aucs / num_events"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "aps = []\n",
    "aucs = []\n",
    "tput = []\n",
    "\n",
    "fig, ax = plt.subplots(1, 2, figsize=(12, 5))\n",
    "\n",
    "for epoch in range(1, EPOCHS + 1):\n",
    "    t0 = time.time()\n",
    "    loss = run_train(tgn_train, train_dl, do_reset=(epoch > 1))\n",
    "    duration = time.time() - t0\n",
    "    tput.append(dataset_size / duration)\n",
    "\n",
    "    aps_epoch, aucs_epoch = run_test(tgn_eval, test_dl)\n",
    "    aps.append(aps_epoch)\n",
    "    aucs.append(aucs_epoch)\n",
    "\n",
    "    ax[0].cla()\n",
    "    ax[0].plot(np.arange(1, epoch + 1), aps)\n",
    "    ax[0].set_xlim(0, EPOCHS)\n",
    "    ax[0].set_ylim(0.95, 0.99)\n",
    "    ax[0].set_xlabel(\"Epoch\")\n",
    "    ax[0].set_ylabel(\"Validation Accuracy\")\n",
    "    ax[1].cla()\n",
    "    ax[1].plot(np.arange(1, epoch + 1), tput)\n",
    "    ax[1].set_xlim(0, EPOCHS)\n",
    "    ax[1].set_xlabel(\"Epoch\")\n",
    "    ax[1].set_ylabel(\"Throughput (samples / s)\")\n",
    "    display.clear_output(wait=True)\n",
    "    display.display(fig)\n",
    "\n",
    "print(f\"Training Finished; Final validation APS {aps_epoch:.4f}, AUCS {aucs_epoch:.4f}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "poptorch_1165",
   "language": "python",
   "name": "poptorch_1165"
  },
  "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.9"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}