Add a PopTorch example (#4)

README.md

@@ -1,6 +1,6 @@
 # Tensor tracker
 
-[API documentation](https://graphcore-research.github.io/pytorch-tensor-tracker/) | [Example](https://graphcore-research.github.io/pytorch-tensor-tracker/usage.html)
+[API documentation](https://graphcore-research.github.io/pytorch-tensor-tracker/) | [Example](doc/Usage.ipynb)
 
 Flexibly track outputs and grad-outputs of `torch.nn.Module`.
 
@@ -37,7 +37,7 @@ display(tracker.to_frame())
 
 <img src="doc/usage_to_frame.png" alt="tensor tracker to_frame output" style="width:30em;"/>
 
-See the [documentation](https://graphcore-research.github.io/pytorch-tensor-tracker/) for more info, or for a more practical example, see our demo of [visualising transformer activations & gradients using UMAP](doc/Example.ipynb).
+See the [documentation](https://graphcore-research.github.io/pytorch-tensor-tracker/) for more info, or for a more practical example, see our demo of [visualising transformer activations & gradients using UMAP](doc/Example.ipynb). To use on IPU with PopTorch, please see [Usage (PopTorch)](doc/UsagePopTorch.ipynb).
 
 
 ## License

doc/UsagePopTorch.ipynb

@@ -0,0 +1,126 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Copyright (c) 2023 Graphcore Ltd. All rights reserved.\n",
+    "\n",
+    "# Usage example (PopTorch)\n",
+    "\n",
+    "Create a toy model to track:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from torch import nn, Tensor\n",
+    "\n",
+    "class Model(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.embed = nn.Embedding(10, 4)\n",
+    "        self.project = nn.Linear(4, 4)\n",
+    "        self.unembed = nn.Linear(4, 10)\n",
+    "\n",
+    "    def forward(self, tokens: Tensor) -> Tensor:\n",
+    "        logits = self.unembed(self.project(self.embed(tokens)))\n",
+    "        return nn.functional.cross_entropy(logits, tokens)\n",
+    "\n",
+    "torch.manual_seed(100)\n",
+    "module = Model()\n",
+    "inputs = torch.randint(0, 10, (3,))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "**PopTorch:**\n",
+    "\n",
+    "A few modifications to work with PopTorch:\n",
+    " - Any tracking should be contained within `forward()`.\n",
+    " - We shouldn't call `tensor.cpu()`, as this is implicit on returned tensors.\n",
+    " - We don't have access to the backward pass."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "[13:48:29.771] [poptorch:cpp] [warning] [DISPATCHER] Type coerced from Long to Int for tensor id 138\n",
+      "Graph compilation: 100%|██████████| 100/100 [00:04<00:00]\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "[Stash(name='embed', type=<class 'torch.nn.modules.sparse.Embedding'>, grad=False, value=tensor([[ 0.4520, -0.1066,  1.1028, -1.1578],\n",
+       "         [-0.4866, -0.1484, -1.6819,  0.7740],\n",
+       "         [-1.0324,  0.2063, -0.7983,  0.4695]])),\n",
+       " Stash(name='project', type=<class 'torch.nn.modules.linear.Linear'>, grad=False, value=tensor([[ 1.2474,  0.4518,  0.2115, -0.6991],\n",
+       "         [-0.3698, -0.1035, -0.2358, -0.3482],\n",
+       "         [ 0.2165,  0.2673, -0.1278, -0.1348]])),\n",
+       " Stash(name='unembed', type=<class 'torch.nn.modules.linear.Linear'>, grad=False, value=tensor([[-0.2676,  0.0945,  0.4727,  0.0716, -0.1146,  0.2311,  0.4380, -0.1172,\n",
+       "           0.6078, -0.0632],\n",
+       "         [ 0.2343, -0.0936,  0.1143, -0.0777,  0.0148, -0.0783,  0.2015,  0.1975,\n",
+       "           0.2441, -0.3956],\n",
+       "         [ 0.1521, -0.0814,  0.2678,  0.0481,  0.1128, -0.0149,  0.3953,  0.2135,\n",
+       "           0.3824, -0.2818]]))]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "from typing import Dict\n",
+    "import poptorch\n",
+    "import tensor_tracker\n",
+    "\n",
+    "class TrackingModel(Model):\n",
+    "    def forward(self, inputs: Tensor) -> Dict[str, Tensor]:\n",
+    "        with tensor_tracker.track(self, stash_value=lambda t: t) as tracker:\n",
+    "            loss = super().forward(inputs)\n",
+    "        return loss, [t.__dict__ for t in tracker]\n",
+    "\n",
+    "loss, tracked = poptorch.inferenceModel(TrackingModel())(inputs)\n",
+    "tracked = [tensor_tracker.Stash(**d) for d in tracked]\n",
+    "display(tracked)\n",
+    "# => [Stash(name=\"embed\", type=nn.Embedding, grad=False, value=tensor(...)),\n",
+    "#     ...]"
+   ]
+  }
+ ],
+ "metadata": {
+  "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.8.10"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}