{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d2b361bf",
   "metadata": {},
   "source": [
    "Copyright (c) 2023 Graphcore Ltd. All rights reserved."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "af14ce5a",
   "metadata": {},
   "source": [
    "# Fraud Detection on IPU using RGCN GNN - Training\n",
    "\n",
    "This notebook demonstrates using PyTorch Geometric on Graphcore IPUs to train a model for fraud detection using the [IEEE-CIS dataset](https://www.kaggle.com/competitions/ieee-fraud-detection/data). The approach is inspired by the [AWS Fraud Detection with GNNs](https://github.com/awslabs/realtime-fraud-detection-with-gnn-on-dgl) project, framing the problem as a node classification task using a heterogeneous graph, where the transaction node types have a label indicating whether they are fraudulent or not.\n",
    "\n",
    "|  Domain | Tasks | Model | Datasets | Workflow |   Number of IPUs   | Execution time |\n",
    "|---------|-------|-------|----------|----------|--------------|--------------|\n",
    "|   GNNs   |  Fraud detection using Node Classification  | RGCN | IEEE-CIS Fraud Detection | Training, evaluation | recommended: 16 (min: 4) | 40 min |\n",
    "\n",
    "In this notebook, you will learn how to:\n",
    "\n",
    "- Take a PyTorch Geometric heterogeneous graph data object, that we loaded in the \"Preprocessing a Tabular Dataset into a PyTorch Geometric Data Object suitable for Fraud Detection\" `1_dataset_preprocessing.ipynb` notebook, and prepare it for training.\n",
    "- Select a model suitable for the task of predicting fraudulent transactions\n",
    "- Train the model on IPUs\n",
    "- Run validation on the trained model\n",
    "\n",
    "This notebook assumes some familiarity with PopTorch as well as PyTorch Geometric (PyG). For additional resources please consult:\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 documentation](https://pytorch-geometric.readthedocs.io/en/latest/)\n",
    "* [PopTorch Geometric documentation](https://docs.graphcore.ai/projects/poptorch-geometric-user-guide/en/latest/index.html)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "db890f52",
   "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-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."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1c89cf77",
   "metadata": {},
   "source": [
    "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."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c2ceb67b",
   "metadata": {},
   "source": [
    "## Dependencies and configuration"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2c881554",
   "metadata": {},
   "source": [
    "Install the dependencies the notebook needs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c72d24a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install  -r requirements.txt\n",
    "%load_ext graphcore_cloud_tools.notebook_logging.gc_logger"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b33a842f",
   "metadata": {},
   "source": [
    "To improve your experience, read in some configuration related to the environment you are running the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "965ca0a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "number_of_ipus = int(os.getenv(\"NUM_AVAILABLE_IPU\", 16))\n",
    "pod_type = os.getenv(\"GRAPHCORE_POD_TYPE\", \"pod16\")\n",
    "executable_cache_dir = os.getenv(\"POPLAR_EXECUTABLE_CACHE_DIR\", \"/tmp/exe_cache/\")\n",
    "dataset_directory = os.getenv(\"DATASETS_DIR\", \".\") + \"/ieee-fraud-detection\"\n",
    "checkpoint_directory = os.getenv(\"CHECKPOINT_DIR\", \".\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0a32779d",
   "metadata": {},
   "source": [
    "## Loading the dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e5823d13",
   "metadata": {},
   "source": [
    "First we will load the dataset for training. We have created a PyTorch Geometric dataset object already containing the preprocessed dataset. If you want to see the preprocessing steps see all the steps in the \"Preprocessing a Tabular Dataset into a PyTorch Geometric Data Object suitable for Fraud Detection\" `1_dataset_preprocessing.ipynb` notebook, or take a look at the dataset object directly in the `dataset.py` script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a97bde6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataset import IeeeFraudDetectionDataset\n",
    "\n",
    "dataset = IeeeFraudDetectionDataset(dataset_directory)\n",
    "data = dataset[0]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0eb6dfba",
   "metadata": {},
   "source": [
    "Let's see what the dataset looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c425d45",
   "metadata": {},
   "outputs": [],
   "source": [
    "data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4f4c90d6",
   "metadata": {},
   "source": [
    "You can see it is a single large heterogeneous graph. The node type we will train on is the `transaction` node, for which we have a label that indicates whether that transaction is fraudulent or not. We have a number of other node types detailing different properties of the transactions, with an edge going out from each transaction to each of the other node types. Full details about the structure of this data can be found in the \"Preprocessing a Tabular Dataset into a PyTorch Geometric Data Object suitable for Fraud Detection\" `1_dataset_preprocessing.ipynb` notebook."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "307a6822",
   "metadata": {},
   "source": [
    "## Preprocessing the dataset\n",
    "\n",
    "Before training on this dataset, we will do some preprocessing.\n",
    "\n",
    "As a first preprocessing step, we apply some transforms on the original graph dataset to:\n",
    "\n",
    " * Make the graph undirected, which will add a reverse edge type for every existing edge type.\n",
    " * Add self loops to all of the node types, which will add a self loop for each of the edge types.\n",
    " * Normalize all the node features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "703d35ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch_geometric.transforms as T\n",
    "\n",
    "data = T.ToUndirected()(data)\n",
    "data = T.AddSelfLoops()(data)\n",
    "data = T.NormalizeFeatures()(data)\n",
    "\n",
    "data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "013f715f",
   "metadata": {},
   "source": [
    "Next, we will create the dataset splits. For this, we will simply use the final 20% of the nodes for validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6da2debf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "num_nodes_train = int(0.8 * data[\"transaction\"].num_nodes)\n",
    "data[\"transaction\"].train_mask = torch.zeros(data[\"transaction\"].num_nodes, dtype=bool)\n",
    "data[\"transaction\"].train_mask[:num_nodes_train] = True\n",
    "data[\"transaction\"].val_mask = torch.zeros(data[\"transaction\"].num_nodes, dtype=bool)\n",
    "data[\"transaction\"].val_mask[num_nodes_train:] = True\n",
    "\n",
    "print(f\"Number of training nodes: {data['transaction'].train_mask.sum()}\")\n",
    "print(f\"Number of validation nodes: {data['transaction'].val_mask.sum()}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "169f84bc",
   "metadata": {},
   "source": [
    "Now let's understand how many transactions in the dataset are actually fraudulent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46c63c85",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_fraud_train = data[\"transaction\"].y[data[\"transaction\"].train_mask].sum()\n",
    "num_total_train = len(data[\"transaction\"].train_mask)\n",
    "num_fraud_val = data[\"transaction\"].y[data[\"transaction\"].val_mask].sum()\n",
    "num_total_val = len(data[\"transaction\"].val_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eab9a5d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of fraud transactions\n",
    "percentage_fraud_train = num_fraud_train / num_total_train\n",
    "percentage_fraud_val = num_fraud_val / num_total_val\n",
    "print(f\"{percentage_fraud_train = :%}\")\n",
    "print(f\"{percentage_fraud_val = :%}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "00c01b9a",
   "metadata": {},
   "source": [
    "We see that there are very few fraudulent transactions in the dataset. As this class imbalance is quite large, we could weight our loss to emphasise fraudulent transactions more than non-fraudulent ones. We can calculate a class weight which we will use later as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3555f6f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use this to set a class weight\n",
    "class_weight = (\n",
    "    (num_total_train / (2 * (num_total_train - num_fraud_train))).item(),\n",
    "    (num_total_train / (2 * num_fraud_train)).item(),\n",
    ")\n",
    "class_weight"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "928fa90d",
   "metadata": {},
   "source": [
    "Later we will see how this class imbalance affects how we track our results. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "dcda1f0b",
   "metadata": {},
   "source": [
    "## Data loading using sampling\n",
    "\n",
    "As the graph we are using is large, we will need some form of sampling to train our model. We will use neighbour sampling. PyTorch Geometric provides a data loader for this, `torch_geometric.loader.NeighborLoader`. Let's create an instance of `NeighborLoader`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46294b42",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64\n",
    "num_layers = 2\n",
    "num_neighbors = [11, 4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "251da6d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.loader import NeighborLoader\n",
    "\n",
    "train_loader = NeighborLoader(\n",
    "    data,\n",
    "    num_neighbors=num_neighbors,\n",
    "    batch_size=batch_size,\n",
    "    input_nodes=(\"transaction\", data[\"transaction\"].train_mask),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6b3951a",
   "metadata": {},
   "outputs": [],
   "source": [
    "next(iter(train_loader))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "679a48e7",
   "metadata": {},
   "source": [
    "When using this data loader, each mini-batch produced has a different shape, depending on how much of the original graph is sampled. The IPU uses ahead-of-time compilation, which requires each mini-batch to have the same size. To achieve this we can use a fixed-size version of this data loader provided in PopTorch Geometric `poptorch_geometric.FixedSizeNeighborLoader`.\n",
    "\n",
    "First, by sampling from the non-fixed-size data loader, we can get an idea of the fixed-sizes we need to accommodate our neighbour loader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86ff2af4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from poptorch_geometric import FixedSizeOptions\n",
    "\n",
    "torch.manual_seed(42)\n",
    "\n",
    "fixed_size_options = FixedSizeOptions.from_loader(train_loader)\n",
    "fixed_size_options"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ff385203",
   "metadata": {},
   "source": [
    "Here we have an approximation for the number of nodes and edges, for each node and edge type, to accommodate the mini-batches produced by the neighbour loader. Now using this with `FixedSizeNeighborLoader`, the mini-batches will be padded to the sizes specified in `fixed_size_options`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59150a4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from poptorch_geometric.neighbor_loader import FixedSizeNeighborLoader\n",
    "\n",
    "train_loader_ipu = FixedSizeNeighborLoader(\n",
    "    data,\n",
    "    num_neighbors=num_neighbors,\n",
    "    fixed_size_options=fixed_size_options,\n",
    "    batch_size=batch_size,\n",
    "    input_nodes=(\"transaction\", data[\"transaction\"].train_mask),\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a7ef9b7c",
   "metadata": {},
   "source": [
    "Now, looking at the first sample you can see it has the dimensions from the `fixed_size_options` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "255e188b",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = next(iter(train_loader_ipu))\n",
    "sample"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3ae4b014",
   "metadata": {},
   "source": [
    "## Picking the right model\n",
    "\n",
    "In order to pick the right model, we should reflect on the task we are doing. For each transaction node we are attempting to predict whether a transaction is fraudulent or not. Each transaction has a number of features as well as being connected to other node types.\n",
    "\n",
    "We will essentially need a [Relational Graph Convolution Network (R-GCN)](https://arxiv.org/abs/1703.06103), where each relation type will have its own set of weights. In this case, each relation type will have its own `SAGEConv` layer. Remember that we only have features for the transactions, so we will need to create some features for the other node types as well. We will train an embedding for each node type for this purpose, after which all the node types will have an embedding to use as an input to the message passing layers.\n",
    "\n",
    "The exact model definition can be seen in `model.py`. Here we:\n",
    " * create the GNN we want to use in homogeneous form\n",
    " * transform the GNN to a heterogeneous GNN, such that we have a convolution layer for each of the edge types\n",
    " * wrap this heterogeneous GNN in a model that contains an embedding for all the non-transaction node types and a loss function such that we can use PopTorch and Graphcore IPUs to train this model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b037f361",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.nn import to_hetero\n",
    "\n",
    "from model import GNN, Model\n",
    "\n",
    "\n",
    "model = GNN(hidden_channels=64, num_layers=num_layers)\n",
    "model = to_hetero(model, data.metadata(), aggr=\"sum\")\n",
    "model = Model(\n",
    "    model,\n",
    "    embedding_size=128,\n",
    "    out_channels=2,\n",
    "    node_types=data.node_types,\n",
    "    num_nodes_per_type={\n",
    "        node_type: data[node_type].num_nodes for node_type in data.node_types\n",
    "    },\n",
    "    class_weight=class_weight,\n",
    ")\n",
    "model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0f497cc4",
   "metadata": {},
   "source": [
    "Using the first sample from the data loader, we can lazily initialize the modules in our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aad2c89f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = next(iter(train_loader_ipu))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f37186e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "out_cpu = model(\n",
    "    sample.x_dict,\n",
    "    sample.edge_index_dict,\n",
    "    batch_size=sample[\"transaction\"].batch_size,\n",
    "    n_id_dict=sample.n_id_dict,\n",
    "    target=sample[\"transaction\"].y,\n",
    "    mask=sample[\"transaction\"].train_mask,\n",
    ")\n",
    "out_cpu"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "14c1fbc5",
   "metadata": {},
   "source": [
    "As a sanity check, we will run inference using this sample on the IPU and verify the results between the IPU and CPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5905665",
   "metadata": {},
   "outputs": [],
   "source": [
    "import poptorch\n",
    "\n",
    "poptorch_options = poptorch.Options()\n",
    "poptorch_options.enableExecutableCaching(executable_cache_dir)\n",
    "inf_model = poptorch.inferenceModel(model, options=poptorch_options)\n",
    "\n",
    "out_ipu = inf_model(\n",
    "    sample.x_dict,\n",
    "    sample.edge_index_dict,\n",
    "    batch_size=sample[\"transaction\"].batch_size,\n",
    "    n_id_dict=sample.n_id_dict,\n",
    "    target=sample[\"transaction\"].y,\n",
    "    mask=sample[\"transaction\"].train_mask,\n",
    ")\n",
    "\n",
    "inf_model.detachFromDevice()\n",
    "\n",
    "out_ipu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c964c29",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert torch.allclose(out_cpu, out_ipu, rtol=1e-05, atol=1e-05)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e5e56511",
   "metadata": {},
   "source": [
    "All looks good."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f4ee6edb",
   "metadata": {},
   "source": [
    "## Training the model\n",
    "\n",
    "We are ready to start training our model. Let's specify the hyperparameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98b652dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 0.01\n",
    "weight_decay = 5e-5\n",
    "num_layers = 2\n",
    "embedding_size = 128\n",
    "hidden_channels = 16\n",
    "log_freq = 10\n",
    "class_weight = (1.00, 5.00)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6d8a0aba",
   "metadata": {},
   "source": [
    "We will train for 50 epochs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "191a985c",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 50"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "76ddcf86",
   "metadata": {},
   "source": [
    "Create the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f68c1798",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = GNN(hidden_channels=hidden_channels, num_layers=num_layers)\n",
    "model = to_hetero(model, data.metadata(), aggr=\"sum\")\n",
    "model = Model(\n",
    "    model,\n",
    "    embedding_size=embedding_size,\n",
    "    out_channels=2,\n",
    "    node_types=data.node_types,\n",
    "    num_nodes_per_type={\n",
    "        node_type: data[node_type].num_nodes for node_type in data.node_types\n",
    "    },\n",
    "    class_weight=class_weight,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "78684a8e",
   "metadata": {},
   "source": [
    "Get the first sample from the data loader to initialize the model lazily."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ee80a18",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader_ipu = FixedSizeNeighborLoader(\n",
    "    data,\n",
    "    num_neighbors=num_neighbors,\n",
    "    fixed_size_options=fixed_size_options,\n",
    "    batch_size=batch_size,\n",
    "    input_nodes=(\"transaction\", data[\"transaction\"].train_mask),\n",
    "    shuffle=True,\n",
    ")\n",
    "\n",
    "sample = next(iter(train_loader_ipu))\n",
    "\n",
    "with torch.no_grad():  # Initialize lazy modules.\n",
    "    out_cpu, loss = model(\n",
    "        sample.x_dict,\n",
    "        sample.edge_index_dict,\n",
    "        batch_size=sample[\"transaction\"].batch_size,\n",
    "        n_id_dict=sample.n_id_dict,\n",
    "        target=sample[\"transaction\"].y,\n",
    "    )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bc1032f9",
   "metadata": {},
   "source": [
    "To accelerate training we will replicate the model over multiple IPUs (4 in this case) and increase the PopTorch `deviceIterations` option to reduce interactions between the host and IPUs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1018781d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import poptorch\n",
    "\n",
    "replication_factor = number_of_ipus\n",
    "device_iterations = 64\n",
    "\n",
    "# Reduce the size of the global batch if it ends up being greater\n",
    "# than the number of training transactions in the dataset\n",
    "if (\n",
    "    data[\"transaction\"].train_mask.sum()\n",
    "    < replication_factor * batch_size * device_iterations\n",
    "):\n",
    "    replication_factor = 1\n",
    "    device_iterations = 1\n",
    "\n",
    "poptorch_options = poptorch.Options()\n",
    "poptorch_options.enableExecutableCaching(executable_cache_dir)\n",
    "poptorch_options.replicationFactor(replication_factor)\n",
    "poptorch_options.deviceIterations(device_iterations)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f38b922a",
   "metadata": {},
   "source": [
    "Re-create the data loader with these options:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2e419f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader_ipu = FixedSizeNeighborLoader(\n",
    "    data,\n",
    "    num_neighbors=num_neighbors,\n",
    "    fixed_size_options=fixed_size_options,\n",
    "    batch_size=batch_size,\n",
    "    input_nodes=(\"transaction\", data[\"transaction\"].train_mask),\n",
    "    shuffle=True,\n",
    "    options=poptorch_options,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b6bc1f40",
   "metadata": {},
   "source": [
    "Wrap the model in `poptorch.trainingModel` specifying the optimizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45cecaf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.train()\n",
    "optimizer = poptorch.optim.Adam(\n",
    "    model.parameters(), lr=learning_rate, weight_decay=weight_decay\n",
    ")\n",
    "training_model = poptorch.trainingModel(\n",
    "    model, optimizer=optimizer, options=poptorch_options\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1da56c71",
   "metadata": {},
   "source": [
    "And finally create the training loop and begin training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9790f281",
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(num_epochs):\n",
    "    total_examples = total_loss = 0\n",
    "    for batch in train_loader_ipu:\n",
    "        out, loss = training_model(\n",
    "            batch.x_dict,\n",
    "            batch.edge_index_dict,\n",
    "            batch_size=sample[\"transaction\"].batch_size,\n",
    "            n_id_dict=batch.n_id_dict,\n",
    "            target=batch[\"transaction\"].y,\n",
    "        )\n",
    "        examples = (\n",
    "            sample[\"transaction\"].batch_size * replication_factor * device_iterations\n",
    "        )\n",
    "        total_examples += examples\n",
    "        total_loss += float(loss.mean()) * examples\n",
    "\n",
    "    if epoch % log_freq == 0:\n",
    "        print(f\"Epoch {epoch}, Loss: {total_loss / total_examples}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "66debd26",
   "metadata": {},
   "source": [
    "Now that the model is trained we can detach it from the IPUs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5940291a",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_model.detachFromDevice()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b107ac22",
   "metadata": {},
   "source": [
    "Let's save the trained weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95a8374e",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(checkpoint_directory, exist_ok=True)\n",
    "torch.save(training_model.state_dict(), os.path.join(checkpoint_directory, \"model.pt\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0c127b86",
   "metadata": {},
   "source": [
    "## Validating the trained model\n",
    "\n",
    "In order to validate the trained model, we must create a data loader that samples from the validation nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89390a04",
   "metadata": {},
   "outputs": [],
   "source": [
    "poptorch_options = poptorch.Options()\n",
    "poptorch_options.enableExecutableCaching(executable_cache_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f50bc91f",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size_val = 1\n",
    "device_iterations = 64\n",
    "replication_factor = number_of_ipus\n",
    "# Reduce the size of the global batch if it ends up being greater\n",
    "# than the number of validation transactions in the dataset\n",
    "if (\n",
    "    data[\"transaction\"].val_mask.sum()\n",
    "    < replication_factor * batch_size_val * device_iterations\n",
    "):\n",
    "    replication_factor = 1\n",
    "    device_iterations = 1\n",
    "\n",
    "poptorch_options.replicationFactor(replication_factor)\n",
    "poptorch_options.deviceIterations(device_iterations)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7b604b2d",
   "metadata": {},
   "source": [
    "For validation, we want to sample the full neighbourhood of the validation nodes. Let's recreate the fixed-size options to ensure we allocate enough space for this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "303382d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_neighbors = [-1, 100]\n",
    "\n",
    "val_loader = NeighborLoader(\n",
    "    data,\n",
    "    num_neighbors=num_neighbors,\n",
    "    batch_size=batch_size_val,\n",
    "    input_nodes=(\"transaction\", data[\"transaction\"].val_mask),\n",
    ")\n",
    "\n",
    "fixed_size_options = FixedSizeOptions.from_loader(val_loader, sample_limit=100)\n",
    "fixed_size_options"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5649d039",
   "metadata": {},
   "source": [
    "And create a fixed-size neighbour loader for validation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa137286",
   "metadata": {},
   "outputs": [],
   "source": [
    "val_loader_ipu = FixedSizeNeighborLoader(\n",
    "    data,\n",
    "    num_neighbors=num_neighbors,\n",
    "    fixed_size_options=fixed_size_options,\n",
    "    batch_size=batch_size_val,\n",
    "    input_nodes=(\"transaction\", data[\"transaction\"].val_mask),\n",
    "    options=poptorch_options,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0544e15d",
   "metadata": {},
   "source": [
    "Let's now wrap the trained model in `poptorch.inferenceModel` and run a single epoch. We take the first `batch_size` number of outputs as this will contain the validation nodes that make up the batch size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "195ed8e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "inference_model = poptorch.inferenceModel(model, options=poptorch_options)\n",
    "\n",
    "outs = []\n",
    "labels = []\n",
    "\n",
    "for batch in val_loader_ipu:\n",
    "    out = inference_model(\n",
    "        batch.x_dict,\n",
    "        batch.edge_index_dict,\n",
    "        batch_size=sample[\"transaction\"].batch_size,\n",
    "        n_id_dict=batch.n_id_dict,\n",
    "    )\n",
    "    outs.append(out[0 :: fixed_size_options.num_nodes[\"transaction\"]])\n",
    "    labels.append(\n",
    "        batch[\"transaction\"].y[0 :: fixed_size_options.num_nodes[\"transaction\"]]\n",
    "    )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e0b1aea2",
   "metadata": {},
   "source": [
    "Again, we will detach the model from IPUs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e10a2253",
   "metadata": {},
   "outputs": [],
   "source": [
    "inference_model.detachFromDevice()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f1be3c3e",
   "metadata": {},
   "source": [
    "## Analysing the results"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0fb51d9a",
   "metadata": {},
   "source": [
    "In this section we will attempt to understand how our trained model performs on the validation nodes.\n",
    "\n",
    "First, we will flatten the results of the validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4e9b1e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = torch.stack(outs)\n",
    "result = result.flatten(start_dim=0, end_dim=1)\n",
    "result.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2bc34ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_true = torch.stack(labels)\n",
    "y_true = y_true.flatten(start_dim=0, end_dim=1)\n",
    "y_true.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "391404e7",
   "metadata": {},
   "source": [
    "We can make our predictions using a softmax function and checking if the second class probability (the fraudulent class) is greater than 0.5."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50bae0c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "y_pred = nn.Softmax(dim=-1)(result)\n",
    "y_pred = y_pred[:, -1]\n",
    "y_pred = y_pred > 0.5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "fde48c2b",
   "metadata": {},
   "source": [
    "And we can get the accuracy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc6d54c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def accuracy(y_pred, y_true):\n",
    "    correct = y_pred.eq(y_true).sum()\n",
    "    return correct / len(y_pred)\n",
    "\n",
    "\n",
    "accuracy(y_pred, y_true)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cb2b25df",
   "metadata": {},
   "source": [
    "Unfortunately, accuracy is a poor metric for this problem as we have such a large class imbalance. Let's instead look at the confusion matrix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3abc358b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_confusion_matrix(y_pred, y_true):\n",
    "    y_pred = y_pred.bool()\n",
    "    y_true = y_true.bool()\n",
    "    true_positives = (y_pred * y_true).sum()\n",
    "    false_positives = (y_pred * ~y_true).sum()\n",
    "    true_negatives = (~y_pred * ~y_true).sum()\n",
    "    false_negatives = (~y_pred * y_true).sum()\n",
    "    return true_positives, false_positives, true_negatives, false_negatives\n",
    "\n",
    "\n",
    "true_pos, false_pos, true_neg, false_neg = get_confusion_matrix(y_pred, y_true)\n",
    "true_pos, false_pos, true_neg, false_neg"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0bee458a",
   "metadata": {},
   "source": [
    "From these we get the true positive and false positive rates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a194e87b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_rates(true_pos, false_pos, true_neg, false_neg):\n",
    "    true_pos_rate = true_pos / (true_pos + false_neg)\n",
    "    false_pos_rate = false_pos / (false_pos + true_neg)\n",
    "    return true_pos_rate, false_pos_rate\n",
    "\n",
    "\n",
    "get_rates(true_pos, false_pos, true_neg, false_neg)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a6f275d7",
   "metadata": {},
   "source": [
    "Now by sweeping over the threshold used to deem a transaction fraudulent, we can get a ROC (receiver operating characteristic) curve."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f646b75",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "results = []\n",
    "for threshold in np.arange(1.0, -0.1, -0.1):\n",
    "    y_pred = nn.Softmax(dim=-1)(result)\n",
    "    y_pred = y_pred[:, -1]\n",
    "    y_pred = y_pred > threshold\n",
    "    results.append((threshold, *get_rates(*get_confusion_matrix(y_pred, y_true))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "deff7167",
   "metadata": {},
   "outputs": [],
   "source": [
    "true_pos_rates = list(zip(*results))[1]\n",
    "false_pos_rates = list(zip(*results))[2]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "09ca878e",
   "metadata": {},
   "source": [
    "And plotting the ROC curve:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d027667c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(false_pos_rates, true_pos_rates)\n",
    "ax.set_xlabel(\"False positive rate\")\n",
    "ax.set_ylabel(\"True positive rate\")\n",
    "plt.grid(True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "154fb92d",
   "metadata": {},
   "source": [
    "A good metric is the area under this curve, let's calculate that:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3adedb27",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "aoc = np.trapz(y=true_pos_rates, x=false_pos_rates)\n",
    "aoc"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ac9e047c",
   "metadata": {},
   "source": [
    "This result is ok for a start, but could do with some improvement. As an extension you could try changing:\n",
    " * the model's layers - perhaps `SAGEConv` layers aren't the best for this use case,\n",
    " * the hyperparameters - does it help to train for more epochs, or maybe even less,\n",
    " * the class weight - are we putting enough weighting on the fraudulent nodes,\n",
    " * the dataset preprocessing - perhaps some features are more useful than others."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7ece4041",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "In this notebook we have seen how to train a heterogeneous GNN model using PyTorch Geometric on Graphcore IPUs for a fraud detection task.\n",
    "\n",
    "Specifically we have:\n",
    "\n",
    " - loaded the preprocessed PyTorch Geometric dataset,\n",
    " - done some additional preprocessing and generated the training and validation splits,\n",
    " - trained a model on the data using neighbour sampling,\n",
    " - validated our trained model by looking at the area under the ROC curve.\n",
    "\n",
    "If you are interested in finding out more about this application check out the \"Preprocessing a Tabular Dataset into a PyTorch Geometric Data Object suitable for Fraud Detection\" `1_dataset_preprocessing.ipynb` notebook. To find out more about PyTorch Geometric on IPUs in general see our PyG tutorials found within `learning-pytorch-geometric-on-ipus/`"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}