{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "61060150-7882-4839-b7c2-6e7bac338266",
   "metadata": {},
   "source": [
    "Copyright (c) 2022 Graphcore Ltd. All rights reserved."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c820e409-be83-40ba-b93e-83f7d82d1227",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Training Dynamic Graphs on IPU using Temporal Graph Networks (TGN)\n",
    "\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",
    "|  Domain | Tasks | Model | Datasets | Workflow |   Number of IPUs   | Execution time |\n",
    "|---------|-------|-------|----------|----------|--------------------|----------------|\n",
    "|   GNNs   |  Link Prediction  | TGN | JODIE | Training, evaluation | recommended: 4 | 20mn |\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."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "392124b0-aa7e-4057-a56c-3e9954c8da42",
   "metadata": {},
   "source": [
    "## Environment setup\n",
    "\n",
    "The best way to run this demo is on Paperspace Gradient’s cloud IPUs because everything is already set up for you. 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/graphcore/Gradient-PyTorch/issues).\n",
    "\n",
    "\n",
    "To run the demo using other IPU hardware, you need to have the Poplar SDK enabled. Refer to the [Getting Started guide](https://docs.graphcore.ai/en/latest/getting-started.html#getting-started) for your system for details on how to enable the Poplar SDK. Also refer to the [Jupyter Quick Start guide](https://docs.graphcore.ai/projects/jupyter-notebook-quick-start/en/latest/index.html) for how to set up Jupyter to be able to run this notebook on a remote IPU machine.\n",
    "\n",
    "Requirements:\n",
    "\n",
    "* Python packages installed with `pip install -r requirements.txt`"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "26bf296d",
   "metadata": {},
   "source": [
    "\n",
    "In order to improve usability and support for future users, Graphcore would like to collect information about the\n",
    "applications and code being run in this notebook. The following information will be anonymised before being sent to Graphcore:\n",
    "\n",
    "- User progression through the notebook\n",
    "- Notebook details: number of cells, code being run and the output of the cells\n",
    "- Environment details\n",
    "\n",
    "You can disable logging at any time by running `%unload_ext graphcore_cloud_tools.notebook_logging.gc_logger` from any cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d530da01",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -r requirements.txt\n",
    "%load_ext graphcore_cloud_tools.notebook_logging.gc_logger"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5c0c45b4",
   "metadata": {},
   "source": [
    "For compatibility with the Paperspace environment variables we will do the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8da4f5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "executable_cache_dir = (\n",
    "    os.getenv(\"POPLAR_EXECUTABLE_CACHE_DIR\", \"/tmp/exe_cache\") + \"/tgn\"\n",
    ")\n",
    "dataset_directory = os.getenv(\"DATASETS_DIR\", \"data\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a5953602",
   "metadata": {},
   "source": [
    "Now we are ready to start!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d6241c7b-bfd4-49dd-992a-a9216e90d1c6",
   "metadata": {},
   "source": [
    "## Load required dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6932104-afb3-4dc8-b46b-648402c0cd24",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os.path as osp\n",
    "\n",
    "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"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "85de4535-84e1-41b1-b1d7-1ddca75d3a43",
   "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."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ed51e2af-1303-4d68-92a9-87358635b62a",
   "metadata": {},
   "source": [
    "#### Define hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7333e5ff-0378-442d-8059-f0b0fcf21e12",
   "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"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "abf2e201-4a0e-4566-a50e-028b77841ea5",
   "metadata": {},
   "source": [
    "#### Create and initialise the TGN model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09b23501-d84f-486c-9bbe-7b649e68e426",
   "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);"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bad1d390-383b-4000-974c-14325ba0c147",
   "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": null,
   "id": "2c299337-9993-48a9-a96f-1476c957f686",
   "metadata": {},
   "outputs": [],
   "source": [
    "jodie_root = osp.join(dataset_directory, \"JODIE\")\n",
    "\n",
    "train_data = DataWrapper(\n",
    "    Data(jodie_root, torch.float32, BATCH_SIZE, NODES_SIZE, EDGES_SIZE), \"train\"\n",
    ")\n",
    "test_data = DataWrapper(\n",
    "    Data(jodie_root, torch.float32, BATCH_SIZE, NODES_SIZE, EDGES_SIZE), \"val\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33b0796e-404b-48f2-96a5-86c340d89421",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_opts = poptorch.Options()\n",
    "train_opts.deviceIterations(DEVICE_ITERATIONS)\n",
    "train_opts.enableExecutableCaching(executable_cache_dir)\n",
    "\n",
    "test_opts = poptorch.Options()\n",
    "test_opts.deviceIterations(1)\n",
    "test_opts.enableExecutableCaching(executable_cache_dir)\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)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c0ec85f3-d5ea-4f5c-887a-4d7e52c894c6",
   "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": null,
   "id": "486779c4-3150-4952-8a21-8c1ca13fbcb9",
   "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)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "31e0d8c8-0d4d-4fab-93f9-9ae733d0e8b7",
   "metadata": {},
   "source": [
    "## Define a training and inference loop and train TGN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cb02a32-9939-4681-9837-0bf2dbcfc0ce",
   "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": null,
   "id": "5749c055-a036-4d5b-b338-339f7d3046f2",
   "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": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}