{ "cells": [ { "attachments": {}, "cell_type": "markdown", "id": "9f8974b8", "metadata": {}, "source": [ "Copyright (c) 2023 Graphcore Ltd. All rights reserved." ] }, { "attachments": {}, "cell_type": "markdown", "id": "d5193627", "metadata": {}, "source": [ "# Molecular Property Prediction on IPU using GIN - Training\n", "\n", "In this tutorial we show an example of how to perform the task of graph classification using the [Graph Isomorphism Network](https://arxiv.org/pdf/1810.00826.pdf) architecture on the Graphcore IPU. This tutorial builds on the [Graph Classification](https://colab.research.google.com/drive/1I8a0DfQ3fI7Njc62__mVXUlcAleUclnb?usp=sharing#scrollTo=1tBMhOrq4JKw) introduction notebook provided by PyTorch Geometric with modifications required in order to use Poptorch. \n", "\n", "| Domain | Tasks | Model | Datasets | Workflow | Number of IPUs | Execution time |\n", "|---------|-------|-------|----------|----------|--------------------|----------------|\n", "| GNNs | Graph Classification | GIN | NCI1 | Training, evaluation | recommended: 4 | ~2 minutes |\n", "\n", "This notebook assumes some familiarity with PopTorch as well as PyTorch Geometric (PyG). For additional resources please consult:\n", "\n", "* [PopTorch Documentation](https://docs.graphcore.ai/projects/poptorch-user-guide/en/latest/index.html)\n", "* [PopTorch Examples and Tutorials](https://docs.graphcore.ai/en/latest/examples.html#pytorch)\n", "* [PyTorch Geometric](https://pytorch-geometric.readthedocs.io/en/latest/)\n", "* [PopTorch Geometric Documentation](https://docs.graphcore.ai/projects/poptorch-geometric-user-guide/en/latest/index.html)\n", "\n", "## 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-Geometric/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": "f859194e", "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": "2ec0e1c5", "metadata": {}, "outputs": [], "source": [ "%pip install -r requirements.txt\n", "%load_ext graphcore_cloud_tools.notebook_logging.gc_logger" ] }, { "attachments": {}, "cell_type": "markdown", "id": "41b326f1", "metadata": {}, "source": [ "And for compatibility with the Paperspace environment variables we will do the following:" ] }, { "cell_type": "code", "execution_count": null, "id": "c6b47fb2", "metadata": {}, "outputs": [], "source": [ "import os\n", "import os.path as osp\n", "\n", "import poptorch\n", "\n", "poptorch.setLogLevel(\"ERR\")\n", "executable_cache_dir = (\n", " os.getenv(\"POPLAR_EXECUTABLE_CACHE_DIR\", \"/tmp/exe_cache/\") + \"/pyg-gin\"\n", ")\n", "dataset_directory = os.getenv(\"DATASETS_DIR\", \"data\")" ] }, { "attachments": {}, "cell_type": "markdown", "id": "8515a9a4", "metadata": {}, "source": [ "Now we are ready to start!" ] }, { "attachments": {}, "cell_type": "markdown", "id": "dc6de011", "metadata": {}, "source": [ "## NCI1 Dataset\n", "\n", "Graph classification refers to the problem of classifying entire graphs (in contrast to nodes), given a dataset of graphs, based on some structural graph properties. The aim is to embed entire graphs in such a way so that they are linearly separable given a task at hand.\n", "\n", "The most common task for graph classification is molecular property prediction, in which molecules are represented as graphs, and the task may be to e.g. infer whether a molecule inhibits HIV virus replication or not.\n", "\n", "The TU Dortmund University has collected a wide range of different graph classification datasets, known as the [TUDatasets](https://chrsmrrs.github.io/datasets/), which are accessible via [torch_geometric.datasets.TUDataset](https://pytorch-geometric.readthedocs.io/en/latest/modules/datasets.html#torch_geometric.datasets.TUDataset) in PyTorch Geometric. Below a medium sized dataset, the NCI1 dataset.\n", "\n", "The NCI1 dataset comes from the cheminformatics domain:\n", "- Each graph in the dataset is molecule that has either a positive or negative association with cell lung cancer\n", "- Each node or vertex in the molecule is an atom and has a node attribute corresponding to the atom's type, represented via a one-hot encoding scheme\n", "- Each edge in the graph corresponds to bonds between the atoms\n", "\n", "Lets first load the dataset:" ] }, { "cell_type": "code", "execution_count": null, "id": "21c4adbc", "metadata": {}, "outputs": [], "source": [ "from torch_geometric.datasets import TUDataset\n", "\n", "tudataset_root = osp.join(dataset_directory, \"TUDataset\")\n", "dataset = TUDataset(root=tudataset_root, name=\"NCI1\")" ] }, { "attachments": {}, "cell_type": "markdown", "id": "bb151abf", "metadata": {}, "source": [ "And take a look at the first item:" ] }, { "cell_type": "code", "execution_count": null, "id": "72633281", "metadata": {}, "outputs": [], "source": [ "data = dataset[0]\n", "print(f\"Graph zero summary:\")\n", "print(f\"Average node degree: {data.num_edges / data.num_nodes:.2f}\")\n", "print(f\"Has isolated nodes: {data.has_isolated_nodes()}\")\n", "print(f\"Has self-loops: {data.has_self_loops()}\")\n", "print(f\"Has edge features: {data.edge_attr is not None}\")\n", "print(f\"Is undirected: {data.is_undirected()}\")\n", "print(f\"Molecule label: {data.y.item()}\")" ] }, { "attachments": {}, "cell_type": "markdown", "id": "e97bf50e", "metadata": {}, "source": [ "We can also get a summary of the entire dataset:" ] }, { "cell_type": "code", "execution_count": null, "id": "df851e44", "metadata": {}, "outputs": [], "source": [ "dataset_summary = dataset.get_summary()\n", "print(dataset_summary)" ] }, { "attachments": {}, "cell_type": "markdown", "id": "37fb147a", "metadata": {}, "source": [ "As expected the mean number of nodes and edges in the dataset is small. We will need to create items of these batches for training, in the next section we will see how to do that." ] }, { "attachments": {}, "cell_type": "markdown", "id": "81b08b86", "metadata": {}, "source": [ "## Data loading for the IPU \n", "\n", "Now that we understand the dataset, lets prepare our data ready for loading.\n", "\n", "First lets split our dataset so we have 80% of the data for training and 20% for evaluation. " ] }, { "cell_type": "code", "execution_count": null, "id": "bba3cb5e", "metadata": {}, "outputs": [], "source": [ "split = int(len(dataset) * 0.8)\n", "train_dataset = dataset[split:]\n", "test_dataset = dataset[:split]" ] }, { "attachments": {}, "cell_type": "markdown", "id": "47653601", "metadata": {}, "source": [ "To provide a batch of graphs to the model rather than individual molecules, we construct dataloaders for both the training and test data. As our data is many small graphs we can use packing to batch our graphs into batches of fixed size, required for the IPU. Packing will make our batches have less wasted space and therefore our computation to be more efficient. For more information on the details of packing see our [Small Graph batching using Packing tutorial](../../../tutorials/tutorials/pytorch_geometric/4_small_graph_batching_with_packing/).\n", "\n", "To do this we need to determine the output size of our batches. First lets start with a batch size of `64` and calculate the sizes to make up the batches:\n", "\n", "- $\\operatorname{MAX}(\\{\\operatorname{NUMNODES}(G_i) ; i=1,\\ldots,N\\}) \\times \\text{batch_size}$\n", "- $\\operatorname{MAX}(\\{\\operatorname{NUMEDGES}(G_i) ; i=1,\\ldots,N\\}) \\times \\text{batch_size}$\n", "\n", "We can use `FixedSizeOptions.from_dataset` to automatically calculate this for this dataset." ] }, { "cell_type": "code", "execution_count": null, "id": "122fd9e2", "metadata": {}, "outputs": [], "source": [ "from poptorch_geometric.fixed_size_options import FixedSizeOptions\n", "\n", "num_graphs = 64\n", "fixed_size_options = FixedSizeOptions.from_dataset(train_dataset, num_graphs)\n", "print(fixed_size_options)" ] }, { "attachments": {}, "cell_type": "markdown", "id": "1fae2659", "metadata": {}, "source": [ "Lets create a dataloader with these. We will also include `add_masks_to_batch` so that the dataloader creates masks that will be easy for us to use in the model." ] }, { "cell_type": "code", "execution_count": null, "id": "ce97a4a9", "metadata": {}, "outputs": [], "source": [ "from poptorch_geometric.dataloader import FixedSizeDataLoader\n", "\n", "train_loader = FixedSizeDataLoader(\n", " dataset=train_dataset,\n", " batch_size=num_graphs,\n", " fixed_size_options=fixed_size_options,\n", " drop_last=True,\n", ")" ] }, { "attachments": {}, "cell_type": "markdown", "id": "bfba984d", "metadata": {}, "source": [ "Lets take a look at the first two samples:" ] }, { "cell_type": "code", "execution_count": null, "id": "c7f03c13", "metadata": {}, "outputs": [], "source": [ "train_loader_iter = iter(train_loader)\n", "first_sample = next(train_loader_iter)\n", "second_sample = next(train_loader_iter)\n", "print(f\"first_sample = {first_sample}\")\n", "print(f\"second_sample = {second_sample}\")" ] }, { "attachments": {}, "cell_type": "markdown", "id": "d585920e", "metadata": {}, "source": [ "Now our batches are fixed size but how much of our batches are wasted space due to padding? Nodes mask can show us how many of our nodes are padded and how many are real." ] }, { "cell_type": "code", "execution_count": null, "id": "3faaef1d", "metadata": {}, "outputs": [], "source": [ "first_sample.nodes_mask" ] }, { "attachments": {}, "cell_type": "markdown", "id": "d475a694", "metadata": {}, "source": [ "And in this first batch how many are real and how many are padded nodes:" ] }, { "cell_type": "code", "execution_count": null, "id": "c60ccd04", "metadata": {}, "outputs": [], "source": [ "float(first_sample.nodes_mask.sum() / len(first_sample.nodes_mask))" ] }, { "attachments": {}, "cell_type": "markdown", "id": "d8bd47ad", "metadata": {}, "source": [ "That is quite a low amount of real compared to padded nodes in this batch, lets try to improve it!\n", "\n", "By default the `FixedSizeDataLoader` uses the `PadToMax` fixed size strategy. This strategy creates mini-batches of a set number of samples and then pads each mini-batch to a fixed-size. Using this strategy enough space must be allocated to ensure that every sampled mini-batch can fit in. This can waste a lot of space in the mini-batch as we some mini-batches may contain many small samples, resulting in a large amount of padding to get the mini-batch up to fixed size. Let's try the `StreamPack` strategy which instead will pack samples into a batch up to the specified fixed-size, resulting in a varying number of samples in a mini-batch but reducing the amount of wasted space. We will also increase the maximum number of graphs in the mini-batch to `300` to allow more graphs to be packed into our mini-batches if there is still space." ] }, { "cell_type": "code", "execution_count": null, "id": "120d3942", "metadata": {}, "outputs": [], "source": [ "from poptorch_geometric import FixedSizeStrategy\n", "\n", "fixed_size_options.num_graphs = 300\n", "\n", "train_loader = FixedSizeDataLoader(\n", " dataset=train_dataset,\n", " batch_size=fixed_size_options.num_graphs,\n", " fixed_size_options=fixed_size_options,\n", " fixed_size_strategy=FixedSizeStrategy.StreamPack,\n", " drop_last=True,\n", ")\n", "first_sample = next(iter(train_loader))\n", "float(first_sample.nodes_mask.sum() / len(first_sample.nodes_mask))" ] }, { "attachments": {}, "cell_type": "markdown", "id": "21912d33", "metadata": {}, "source": [ "That's much better, our batches should be much more efficient and if you inspect the sizes they will still be the sizes we obtained before. One thing we will have to bear in mind going forward is that our batches may not all have the same number of real graphs in as padded graphs. We can see that by looking at the `graphs_mask` for each batch:" ] }, { "cell_type": "code", "execution_count": null, "id": "82ef9843", "metadata": {}, "outputs": [], "source": [ "for i, batch in enumerate(train_loader):\n", " print(f\"Batch {i} has {batch.graphs_mask.sum()} real graphs\")" ] }, { "attachments": {}, "cell_type": "markdown", "id": "03c98666", "metadata": {}, "source": [ "We can also configure the above dataloader to improve IO performance. This comes free from the features supported by the PopTorch dataloader. For more information on efficient data batching/loading, please refer to the resources below:\n", "- PopTorch documentation [Efficient data batching](https://docs.graphcore.ai/projects/poptorch-user-guide/en/latest/batching.html#efficient-data-batching)\n", "- PopTorch tutorial: [Efficient data loading](https://github.com/graphcore/tutorials/tree/sdk-release-2.5/tutorials/pytorch/tut2_efficient_data_loading)\n", "\n", "One key thing to note is that when increasing the device iterations it is useful to also change the output mode, so that the output for each iteration is returned after each step.\n", "\n", "We also enable caching the executable so we don't need to recompile our model if we have already." ] }, { "cell_type": "code", "execution_count": null, "id": "7e41b8cf", "metadata": {}, "outputs": [], "source": [ "poptorch_options = poptorch.Options()\n", "poptorch_options.deviceIterations(2)\n", "poptorch_options.outputMode(poptorch.OutputMode.All)\n", "poptorch_options.enableExecutableCaching(executable_cache_dir)" ] }, { "attachments": {}, "cell_type": "markdown", "id": "1f4f45ce", "metadata": {}, "source": [ "Then recreate our dataloader with these options:" ] }, { "cell_type": "code", "execution_count": null, "id": "6c286dc3", "metadata": {}, "outputs": [], "source": [ "train_loader = FixedSizeDataLoader(\n", " dataset=train_dataset,\n", " batch_size=fixed_size_options.num_graphs,\n", " fixed_size_options=fixed_size_options,\n", " fixed_size_strategy=FixedSizeStrategy.StreamPack,\n", " drop_last=True,\n", " options=poptorch_options,\n", ")" ] }, { "attachments": {}, "cell_type": "markdown", "id": "b8fa191e", "metadata": {}, "source": [ "Taking a look at the first batch you can see the items are twice as big as you would expect. This is because we have chosen two device iterations. PopTorch will split this single large batch between the two iterations." ] }, { "cell_type": "code", "execution_count": null, "id": "9a20b0a2", "metadata": {}, "outputs": [], "source": [ "next(iter(train_loader))" ] }, { "attachments": {}, "cell_type": "markdown", "id": "558d743a", "metadata": {}, "source": [ "We have a dataloader that is ready to feed data into a model, in the next section we will take a look at the model itself." ] }, { "attachments": {}, "cell_type": "markdown", "id": "962ba9c8", "metadata": {}, "source": [ "## Graph Isomorphism Network Architecture\n", "\n", "Generally, GNNs designed for graph classification build on top of architectures for node **classification** by adding a **readout** layer and classification layer at the end of the network:\n", "\n", "- The readout layer is used to aggregate node embeddings to generate an embedding for each graph\n", "- The classification layer is used to accomplish the task of graph classification\n", "\n", "Graph node embeddings using GNN architectures are typically characterised via an _aggregation_ phase and an _update_ or _combine_ phase to give a nodes $k$-hop representation:\n", "\n", "$$a_v^{(k)}=\\operatorname{AGGREGATE}^{(k)}\\left(\\left\\{h_u^{(k-1)}: u \\in \\mathcal{N}(v)\\right\\}\\right), \\quad h_v^{(k)}=\\operatorname{COMBINE}^{(k)}\\left(h_v^{(k-1)}, a_v^{(k)}\\right)$$\n", "\n", "The Graph Isomorphism Network specialises this scheme by defining:\n", "\n", "$$ a_v^{(k)} := \\sum_{u\\in \\mathcal{N}{(v)}} h_u^{(k-1)},~~ h_{v}^{(k)} := \\operatorname{MLP}^{(k)}((1+\\epsilon^{(k)}) \\cdot h_{v}^{(k-1)} + a_v^{(k)})$$\n", "\n", "Where $h_{v}^{(0)}$ are the initial node attributes for each node. The above specialisation is a result of wanting to represent an injective function that operates on multisets of node attributes through the use of MLPs. If the $k$-hop representation of a node is defined by an injective function, it can be asserted that any two nodes that have the same $k$-hop representation are isomorphic due to the definition of injectivity. Further, any two graphs that have the same representation as generated by such an injective function will be isomorphic. This property is well suited for the task of graph classification as it allows graphs with different structural properties to differentiated. Some GNN architectures (e.g. [GraphSage](https://arxiv.org/abs/1706.02216), [GCN](https://arxiv.org/abs/1609.02907)) do not satisfy learning injective functions on multisets and therefore have drawbacks when used for the graph classification task.\n", "\n", "To generate graph level embeddings, GIN defines the below **readout** layer which concatenates graph representations obtained from each $k$-hop node representation:\n", "\n", "$$G_i = \\operatorname{CONCAT}(\\operatorname{LINEAR}(\\sum_{v ;~ v \\in G_i } h_v^{(k)}) ~|~ k=0,\\ldots,K)$$\n", "\n", "Unnormalized logits for each graph can therefore be obtained by applying a subsequent $\\operatorname{LINEAR}$ layer." ] }, { "attachments": {}, "cell_type": "markdown", "id": "c47b9f81", "metadata": {}, "source": [ "Below is an implementation of GIN using PyTorch Geometric and Poptorch. Some clarifying notes and differences:\n", "\n", "- The norm used is layer normalization as it doesn't require modifications when using padding in our batches and provides similar results to batch normalization.\n", "- The loss is required to be computed in the forward method for compatibility with PopTorch.\n", "- The loss is aggregated on all but the final padding graph, see below for more information.\n", "- For **readout** and **classification**, rather than concatenating all $hop-k$ graph layer representations and then applying another $\\operatorname{LINEAR}$ layer to obtain classification scores, the $\\operatorname{LINEAR}$ layer in the definition of $G_i$ above is chosen to output classification scores, which are then aggregated over all hops. This is equivalent to $\\operatorname{CONCAT}$ followed by $\\operatorname{LINEAR}$ as a composition of linear functions is linear, so the learned linear map will be the same. This matches the [original implementation](https://github.com/weihua916/powerful-gnns)\n", "\n", "For simplicity we have put the model implementation in `model.py` and we will import it here. Feel free to look over the implementation and see the points mentioned above." ] }, { "cell_type": "code", "execution_count": null, "id": "8b7075aa", "metadata": {}, "outputs": [], "source": [ "from model import GIN" ] }, { "attachments": {}, "cell_type": "markdown", "id": "13d6fcfc", "metadata": {}, "source": [ "Now we have a good understanding of the model architecture, lets start training the model." ] }, { "attachments": {}, "cell_type": "markdown", "id": "529d105d", "metadata": {}, "source": [ "## Training a GIN for graph classification\n", "\n", "Now we can train and evaluate the GIN on the IPU by using the dataloaders and GIN architecture introduced above. First we set some hyperparameters and configure all that is required in order to execute our model using [PopTorch](https://github.com/graphcore/tutorials/tree/sdk-release-3.0/tutorials/pytorch/basics)." ] }, { "cell_type": "code", "execution_count": null, "id": "847da3b0", "metadata": {}, "outputs": [], "source": [ "in_channels = dataset.num_node_features\n", "hidden_channels = 32\n", "# output hidden units dimension, one unit for each class\n", "out_channels = 2\n", "# number of GinConv layers\n", "num_conv_layers = 4\n", "# number of hidden layers used in the mlp for each GinConv layer\n", "num_mlp_layers = 2\n", "\n", "learning_rate = 1e-2\n", "num_epochs = 200" ] }, { "attachments": {}, "cell_type": "markdown", "id": "938520ec", "metadata": {}, "source": [ "Now we can create the model. We will use a PopTorch optimizer that is functionally the same but more suited to the IPU for performance and memory." ] }, { "cell_type": "code", "execution_count": null, "id": "37517c89", "metadata": {}, "outputs": [], "source": [ "model = GIN(\n", " in_channels=in_channels,\n", " hidden_channels=hidden_channels,\n", " out_channels=out_channels,\n", " num_conv_layers=num_conv_layers,\n", " num_mlp_layers=num_mlp_layers,\n", " batch_size=fixed_size_options.num_graphs,\n", ")\n", "\n", "model.train()\n", "model" ] }, { "attachments": {}, "cell_type": "markdown", "id": "851d6417", "metadata": {}, "source": [ "Now we can make it into a PopTorch model ready for training:" ] }, { "cell_type": "code", "execution_count": null, "id": "f77bfe8f", "metadata": {}, "outputs": [], "source": [ "optimizer = poptorch.optim.Adam(model.parameters(), lr=learning_rate)\n", "\n", "poptorch_training_model = poptorch.trainingModel(\n", " model, optimizer=optimizer, options=poptorch_options\n", ")" ] }, { "attachments": {}, "cell_type": "markdown", "id": "dd047a03", "metadata": {}, "source": [ "Now train the model for `num_epochs`. Note below that we decay the learning rate by a factor of 0.5 every 50 epochs. " ] }, { "cell_type": "code", "execution_count": null, "id": "29eae07d", "metadata": {}, "outputs": [], "source": [ "from tqdm import tqdm\n", "\n", "epoch_bar = tqdm(range(num_epochs))\n", "\n", "epoch_losses = []\n", "for epoch in epoch_bar:\n", " epoch_loss = 0\n", " for data in train_loader:\n", " preds, micro_batch_loss = poptorch_training_model(\n", " data.x,\n", " data.edge_index,\n", " data.batch,\n", " graphs_mask=data.graphs_mask,\n", " target=data.y,\n", " )\n", " epoch_loss += micro_batch_loss.mean()\n", "\n", " # decay learning rate every 50 epochs\n", " if (epoch + 1) % 50 == 0:\n", " learning_rate *= 0.5\n", " optimizer = poptorch.optim.Adam(model.parameters(), lr=learning_rate)\n", " poptorch_training_model.setOptimizer(optimizer)\n", "\n", " epoch_bar.set_description(f\"Epoch {epoch} training loss: {epoch_loss:0.6f}\")\n", "\n", " epoch_losses.append(epoch_loss)" ] }, { "attachments": {}, "cell_type": "markdown", "id": "7e3ebcf8", "metadata": {}, "source": [ "Now training is complete we can detach from the IPU:" ] }, { "cell_type": "code", "execution_count": null, "id": "5cb1f87a", "metadata": {}, "outputs": [], "source": [ "poptorch_training_model.detachFromDevice()" ] }, { "attachments": {}, "cell_type": "markdown", "id": "c403743f", "metadata": {}, "source": [ "Lets take a look at the loss curve:" ] }, { "cell_type": "code", "execution_count": null, "id": "b05e659e", "metadata": {}, "outputs": [], "source": [ "from matplotlib import pyplot as plt\n", "\n", "fig, ax = plt.subplots()\n", "ax.plot(list(range(num_epochs)), epoch_losses)\n", "ax.set_xlabel(\"Epoch\")\n", "ax.set_ylabel(\"Loss\")\n", "plt.grid(True)" ] }, { "attachments": {}, "cell_type": "markdown", "id": "e6aad59d", "metadata": {}, "source": [ "## Evaluation\n", "\n", "Now we have trained our model we can evaluate it on our test dataset. First lets create a test dataloader:" ] }, { "cell_type": "code", "execution_count": null, "id": "b8750e16", "metadata": {}, "outputs": [], "source": [ "test_loader = FixedSizeDataLoader(\n", " dataset=test_dataset,\n", " batch_size=fixed_size_options.num_graphs,\n", " fixed_size_options=fixed_size_options,\n", " fixed_size_strategy=FixedSizeStrategy.StreamPack,\n", " drop_last=True,\n", " options=poptorch_options,\n", ")" ] }, { "attachments": {}, "cell_type": "markdown", "id": "ea5cfa77", "metadata": {}, "source": [ "Similarly, we wrap the model with the PopTorch inferenceModel wrapper in order for the model to execute in evaluation mode and run on the IPU. We must take special care when calculating the accuracy, ensuring we mask out any of the padded graphs before calculating the final accuracy." ] }, { "cell_type": "code", "execution_count": null, "id": "192eacaf", "metadata": {}, "outputs": [], "source": [ "model.eval()\n", "poptorch_inference_model = poptorch.inferenceModel(model, options=poptorch_options)\n", "\n", "total_correct = 0\n", "num_samples = 0\n", "for data in test_loader:\n", " scores = poptorch_inference_model(data.x, data.edge_index, data.batch)\n", " preds = scores.argmax(dim=1)\n", " total_correct += (\n", " (preds.flatten() == data.y.flatten()) * data.graphs_mask.flatten()\n", " ).sum()\n", " num_samples += data.graphs_mask.sum()\n", "\n", "accuracy = total_correct / num_samples\n", "print(f\"Test accuracy: {accuracy}\")" ] }, { "attachments": {}, "cell_type": "markdown", "id": "263b6890", "metadata": {}, "source": [ "Finally, lets detach the inference model from the IPUs:" ] }, { "cell_type": "code", "execution_count": null, "id": "ff6234bf", "metadata": {}, "outputs": [], "source": [ "poptorch_inference_model.detachFromDevice()" ] }, { "attachments": {}, "cell_type": "markdown", "id": "82fb93ba", "metadata": {}, "source": [ "We have successfully trained a GIN model and run evaluation on IPUs using packing to speed up our training." ] }, { "attachments": {}, "cell_type": "markdown", "id": "19100545", "metadata": {}, "source": [ "## Follow up\n", "\n", "Here we have seen how to train a GIN model to predict properties of small graphs. We have used packing to make efficient use of our batches while achieving fixed size.\n", "\n", "Next you can try:\n", "\n", "* Can you achieve a higher by adjusting some of the hyperparameters?\n", "* Take a look at our high performing molecular prediction example using [SchNet](../../schnet/pytorch_geometric/)." ] } ], "metadata": { "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }