ipu/examples/jax_ipu_gnn_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2ae94e2f-edce-4659-b74e-a229c2e44ff7",
   "metadata": {},
   "source": [
    "# JAX on IPU: GNN Tutorial\n",
    "In this tutorial we use JAX on Graphcore IPUs to build a simple GNN for a small node classification task. For educational purposes we rely on plain JAX without using any higher-level libraries such as [Flax](https://github.com/google/flax), [Haiku](https://github.com/deepmind/dm-haiku), or [Jraph](https://github.com/deepmind/jraph)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b252595-f17b-4c62-bfb5-d56328ac5871",
   "metadata": {},
   "source": [
    "We enable a compilation progress bar to see when our model gets (re-)compiled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a35a8a7e",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%env TF_POPLAR_FLAGS=--show_progress_bar=True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3921f459",
   "metadata": {},
   "source": [
    "First install and import some dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2672be91-f754-411d-9a09-19244cb821e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found existing installation: jax 0.3.16+ipu\n",
      "Uninstalling jax-0.3.16+ipu:\n",
      "  Successfully uninstalled jax-0.3.16+ipu\n",
      "Found existing installation: jaxlib 0.3.15+ipu.sdk310\n",
      "Uninstalling jaxlib-0.3.15+ipu.sdk310:\n",
      "  Successfully uninstalled jaxlib-0.3.15+ipu.sdk310\n"
     ]
    }
   ],
   "source": [
    "# Install experimental JAX for IPUs (SDK 3.1) from Github releases.\n",
    "import sys\n",
    "!{sys.executable} -m pip uninstall -y jax jaxlib\n",
    "!{sys.executable} -m pip install -q jax==0.3.16+ipu jaxlib==0.3.15+ipu.sdk310 -f https://graphcore-research.github.io/jax-experimental/wheels.html\n",
    "!{sys.executable} -m pip install -q matplotlib\n",
    "!{sys.executable} -m pip install -q networkx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdaeb1c1-743f-4489-a552-70ccd49df8cd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "from jax.config import config\n",
    "import networkx as nx\n",
    "import matplotlib.pyplot as plt\n",
    "from functools import partial\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6330d24b-8dca-40d4-97ac-f83d089b7a4e",
   "metadata": {},
   "source": [
    "We set `jax_platforms = \"cpu,ipu\"` for using CPU as the default platform for initialization of parameters and dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdeea884-b00a-4882-9be2-bd3734cba00e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "config.FLAGS.jax_platforms = \"cpu,ipu\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f2a58cb-fd74-45d5-957c-45eda640e269",
   "metadata": {},
   "source": [
    "We can switch between devices to train the model by setting the `DEVICE` variable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a4e9727-a91b-4e47-b3f3-3c127cce689d",
   "metadata": {},
   "outputs": [],
   "source": [
    "DEVICE = \"ipu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb7a7dac-ef25-4a72-9b1f-24003fa820aa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "devices = jax.devices(DEVICE)\n",
    "print(devices)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "111a163a-ce25-4b22-aa7d-6b0ae1c2eee1",
   "metadata": {},
   "source": [
    "## Define the graph\n",
    "For this notebook we use [Zachary's karate club](https://en.wikipedia.org/wiki/Zachary%27s_karate_club) graph, a well-known example for node classification on a small social graph: The 34 nodes of the graph represent the members of a karate club, edges represent social interactions between club members. A conflict between the members represented by nodes 0 and 33 lead to a splitting of the club. The task is to predict for every member which of the two new clubs they are going to join."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be60e0ca-0db2-4861-8369-2535e08150f1",
   "metadata": {
    "tags": []
   },
   "source": [
    "We can define the graph in this notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2245279-ee22-4ae8-b1d5-7a3f655199c5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "edges = [(1, 0), (2, 0), (2, 1), (3, 0), (3, 1), \n",
    "         (3, 2), (4, 0), (5, 0), (6, 0), (6, 4), \n",
    "         (6, 5), (7, 0), (7, 1), (7, 2), (7, 3), \n",
    "         (8, 0), (8, 2), (9, 2), (10, 0), (10, 4), \n",
    "         (10, 5), (11, 0), (12, 0), (12, 3), (13, 0), \n",
    "         (13, 1), (13, 2), (13, 3), (16, 5), (16, 6), \n",
    "         (17, 0), (17, 1), (19, 0), (19, 1), (21, 0), \n",
    "         (21, 1), (25, 23), (25, 24), (27, 2), (27, 23), \n",
    "         (27, 24), (28, 2), (29, 23), (29, 26), (30, 1), \n",
    "         (30, 8), (31, 0), (31, 24), (31, 25), (31, 28), \n",
    "         (32, 2), (32, 8), (32, 14), (32, 15), (32, 18), \n",
    "         (32, 20), (32, 22), (32, 23), (32, 29), (32, 30), \n",
    "         (32, 31), (33, 8), (33, 9), (33, 13), (33, 14), \n",
    "         (33, 15), (33, 18), (33, 19), (33, 20), (33, 22), \n",
    "         (33, 23), (33, 26), (33, 27), (33, 28), (33, 29), \n",
    "         (33, 30), (33, 31), (33, 32)]\n",
    "\n",
    "node_labels = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])\n",
    "num_nodes = len(node_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f91bb65-cd65-4539-a59d-14787ca70ea6",
   "metadata": {},
   "source": [
    "#### Visualisation\n",
    "Next we create a networkx graph for visualisation and represent the graph as a `jax.numpy` array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "382f78ae-09d3-4af9-8c28-e38a69c675a2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "g = nx.Graph()\n",
    "g.add_nodes_from(range(num_nodes))\n",
    "g.add_edges_from(edges)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69465c71-ebfc-42a4-8a56-7245e50ac882",
   "metadata": {},
   "source": [
    "Now we can draw the graph of the club members and their social interactions. The members 0 and 33 are treated as the only ones with know label corresponding to their new clubs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fa921c7-ba14-4240-94fe-1c7729ebfc46",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "c_0 = np.array([0.1, 0.5, 1.0])\n",
    "c_1 = np.array([1.0, 0.5, 0.1])\n",
    "c_default = np.array([.5, .5, .5])\n",
    "c_error = np.array([1.0, 0.1, 0.1])\n",
    "\n",
    "color_map = [c_default for _ in range(num_nodes)]\n",
    "color_map[0] = c_0\n",
    "color_map[-1] = c_1\n",
    "\n",
    "NODE_SIZE = 200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a3294ae-63d5-483f-8728-379ef727ebb5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pos = nx.kamada_kawai_layout(g)\n",
    "fig, ax = plt.subplots(1, 1, figsize=[6, 7])\n",
    "nx.draw(g, pos, node_size=NODE_SIZE, node_color=color_map, with_labels=True, font_color=\"w\", font_size=10, ax=ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14d28a8d-7a8d-447e-819b-79aa4e9eb57a",
   "metadata": {},
   "source": [
    "For modelling purposes we add inverse edges as well as self-loops to the graph:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a1dd94b-80c5-4976-964e-5562919a60f3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "all_edges = edges + [(edge[1], edge[0]) for edge in edges] + [(i, i) for i in range(num_nodes)]\n",
    "graph = jnp.array(all_edges)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77a281d5-885e-47b7-b652-4308538e2589",
   "metadata": {},
   "source": [
    "## Model definition\n",
    "We solve this task with a simple Graph Convolutional Network (GCN) ([Kipf, Welling, 2016](https://arxiv.org/abs/1609.02907)).\n",
    "First, we define functions to explicitly initialise the parameters and apply a GCN layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68b8c5a6-6085-4577-8fb6-8de3d1ee8a90",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def gcn_layer(input_size, output_size, nonlinearity=None, use_bias=False):\n",
    "    def parameter_init(key, scale=0.02):\n",
    "        if use_bias:\n",
    "            return (scale * jax.random.normal(key, (output_size, input_size)), jnp.zeros(output_size))\n",
    "        return (scale * jax.random.normal(key, (output_size, input_size)), )\n",
    "    \n",
    "    def apply(params, node_embeddings, graph):\n",
    "        node_embeddings = jnp.dot(params[0], node_embeddings.T).T\n",
    "        if use_bias:\n",
    "            node_embeddings = node_embeddings + params[1]\n",
    "        if nonlinearity:\n",
    "            node_embeddings = nonlinearity(node_embeddings)\n",
    "        messages = node_embeddings[graph[:, 0]]\n",
    "        node_embeddings = jax.ops.segment_sum(messages, graph[:, 1], num_nodes)\n",
    "        return node_embeddings, graph\n",
    "    \n",
    "    return parameter_init, apply"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ee427d7-a804-4890-9812-22a7a3d6e034",
   "metadata": {},
   "source": [
    "Now, we can define a multi-layer GCN in a similar way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c0f75d3-d6ab-46c1-a976-55b05338620b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def gcn(layer_sizes):\n",
    "    layers = []\n",
    "    parameter_inits = []\n",
    "    for n, (input_size, output_size) in enumerate(zip(layer_sizes[:-1], layer_sizes[1:])):\n",
    "        layer_parameter_init, layer_apply = gcn_layer(\n",
    "            input_size,\n",
    "            output_size,\n",
    "            nonlinearity=jax.nn.relu if n < len(layer_sizes) - 2 else None,\n",
    "            use_bias=False\n",
    "        )\n",
    "        parameter_inits.append(layer_parameter_init)\n",
    "        layers.append(layer_apply)\n",
    "    \n",
    "    def parameter_init(key, scale=0.02):\n",
    "        keys = jax.random.split(key, len(layer_sizes))\n",
    "        params = []\n",
    "        for layer_parameter_init, layer_key in zip(parameter_inits, keys):\n",
    "            params.append(layer_parameter_init(layer_key, scale))\n",
    "        return params\n",
    "    \n",
    "    def apply(params, node_embeddings, graph):\n",
    "        for layer, layer_params in zip(layers, params):\n",
    "            node_embeddings, graph = layer(layer_params, node_embeddings, graph)\n",
    "        return node_embeddings\n",
    "    \n",
    "    return parameter_init, apply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f208b8c5-936a-433d-b1c2-0430d6794b9d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "layer_size = [num_nodes] + [64, 64] + [2]\n",
    "gcn_init, gcn_predict = gcn(layer_size)\n",
    "params = gcn_init(jax.random.PRNGKey(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "679ca078-fb74-408a-a7b6-49411b0c6110",
   "metadata": {},
   "source": [
    "Finally, we define a prediction function that returns the probability of a node belonging to the two different classes, a loss function, a validation function and a training step that uses plain SGD to update the model parameters. Note that the loss function only uses the embeddings of the only two nodes whose label we know at this point: node 0 and node 33. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e68a3db5-9ac0-46af-8fbd-6785ba1fe88b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def prediction(params, graph):\n",
    "    initial_node_embeddings = jnp.eye(num_nodes)\n",
    "    return jax.nn.softmax(gcn_predict(params, initial_node_embeddings, graph))\n",
    "\n",
    "def loss_fun(params, graph):\n",
    "    log_probs = jnp.log(prediction(params, graph))\n",
    "    return -(log_probs[0, 0] + log_probs[-1, 1]) / 2\n",
    "\n",
    "def validation_fun(params, graph, node_labels):\n",
    "    probs = prediction(params, graph)\n",
    "    acc = jnp.mean(jnp.argmax(probs, axis=1) == node_labels)\n",
    "    log_probs = jnp.log(probs)\n",
    "    valid_loss = -jnp.mean(log_probs[:,0] * (1 - node_labels) + log_probs[:,1] * node_labels)\n",
    "    return probs, valid_loss, acc\n",
    "\n",
    "# Explicit jitting for IPU backend.\n",
    "# Donate `params`` to keep parameters on IPU SRAM. \n",
    "@partial(jax.jit, backend=DEVICE, donate_argnums=(0,))\n",
    "def training_step(params, graph, node_labels, learning_rate):\n",
    "    loss, grads = jax.value_and_grad(loss_fun)(params, graph)\n",
    "    updated_params = [tuple(p - learning_rate * dp for p, dp in zip(param, d_param)) for param, d_param in zip(params, grads)]\n",
    "    probs, valid_loss, acc = validation_fun(params, graph, node_labels)\n",
    "    stats = {\"probabilities\": probs,\n",
    "             \"training_loss\": loss,\n",
    "             \"validation_loss\": valid_loss,\n",
    "             \"accuracy\": acc}\n",
    "    return stats, updated_params"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c40ffa0e-0608-4fa0-be31-d0262f5e9937",
   "metadata": {},
   "source": [
    "## Training\n",
    "We train the model on IPU for 20 steps. The first step includes compilation and therefore takes longer.\n",
    "\n",
    "Every fifth step we visualise the node probabilities. The accuracy (proportion of correctly classified nodes) should reach a steady state of 0.971 (corresponding to 33 / 34 correctly classified nodes) quite early on while the validation loss still drops, showing a better separation of the two classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "291fa777-7198-4345-9ade-7570bb877d41",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def visualise(prob, ax):\n",
    "    color_map_predicted = [p[0] * c_0 + p[1] * c_1 for p in prob]\n",
    "    color_map_predicted[0] = c_0\n",
    "    color_map_predicted[-1] = c_1\n",
    "    \n",
    "    nx.draw(g, pos, node_size=NODE_SIZE, node_color=color_map_predicted, with_labels=True, font_color=\"w\", font_size=10, ax=ax)\n",
    "    ax.set_title(\"Soft Predictions\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48d04eaf-6c70-4979-ad5f-0ccd93a73713",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "learning_rate = 0.02\n",
    "num_steps = 20\n",
    "vis_step = 5\n",
    "\n",
    "fig, ax = plt.subplots(1, num_steps//vis_step, figsize=[20, 6])\n",
    "fig.suptitle(\"Node classification\", fontsize=20)\n",
    "fig.tight_layout()\n",
    "\n",
    "for step in range(1, num_steps + 1):\n",
    "    t0 = time.time()\n",
    "    stats, params = training_step(params, graph, node_labels, learning_rate)\n",
    "    if step % vis_step == 0:\n",
    "        visualise(np.array(stats[\"probabilities\"]), ax[step // vis_step - 1])\n",
    "        ax[step // vis_step - 1].set_title(f\"Step {step}\")\n",
    "    print(f\"Step {step}, duration = {(time.time() - t0) * 1000:.2f} ms, \"\n",
    "          f\"Validation Loss = {stats['validation_loss']:.3f}, Accuracy = {stats['accuracy']:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "486af8b1-4834-4d51-94f9-30c378164fc2",
   "metadata": {},
   "source": [
    "We now plot the final results. We should find that node 8 gets misclassified, an error observed in many predictions on this dataset, including in [Zachary's original 1977 publication](https://www.jstor.org/stable/3629752)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0a5e228-28a1-4b03-91a9-5a9db8d43fa9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "params_host = jax.device_get(params)\n",
    "probs = np.array(prediction(params_host, graph))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15d0d2e5-a6c8-49b9-8bd6-1f63d5436b10",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "predicted_class = np.argmax(probs, axis=1)\n",
    "errors = (node_labels != predicted_class)\n",
    "\n",
    "color_map_predicted = [(1-p) * 0.75 * c_0 + p * 0.75 * c_1 for p in predicted_class]\n",
    "color_map_predicted[0] = c_0\n",
    "color_map_predicted[-1] = c_1\n",
    "\n",
    "color_map_labels = [l * 0.75 * c_1 + (1 - l) * 0.75 * c_0 for l in node_labels]\n",
    "color_map_labels[0] = c_0\n",
    "color_map_labels[-1] = c_1\n",
    "\n",
    "color_map_err = [l * c_error + (1 - l) * c_default for l in errors]\n",
    "color_map_err[0] = c_0\n",
    "color_map_err[-1] = c_1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b0b87c9-2b0b-408b-a59a-969a3af2d0de",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 3, figsize=[18, 7])\n",
    "nx.draw(g, pos, node_size=NODE_SIZE, node_color=color_map_predicted, with_labels=True, font_color=\"w\", font_size=10, ax=ax[0])\n",
    "nx.draw(g, pos, node_size=NODE_SIZE, node_color=color_map_labels, with_labels=True, font_color=\"w\", font_size=10, ax=ax[1])\n",
    "nx.draw(g, pos, node_size=NODE_SIZE, node_color=color_map_err, with_labels=True, font_color=\"w\", font_size=10, ax=ax[2])\n",
    "\n",
    "ax[0].set_title(\"Predicted Labels\")\n",
    "ax[1].set_title(\"Ground Truth\")\n",
    "ax[2].set_title(\"Difference\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab5f2634-47bd-4388-b869-7952823e2528",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}