{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a7acdeab",
   "metadata": {},
   "source": [
    "# Multi-Horizon Financial Forecasting on IPU using DeepLOB-Attention - Training\n",
    "\n",
    "The [original Jupyter notebook](https://github.com/zcakhaa/Multi-Horizon-Forecasting-for-Limit-Order-Books/blob/main/code_gpu/run_train_deeplob_attention.ipynb) was based on the paper: [\"Multi-Horizon Forecasting for Limit Order Books: Novel Deep Learning Approaches and Hardware Acceleration using Intelligent Processing Units\"](https://arxiv.org/abs/2105.10430).\n",
    "Original authors: Zihao Zhang and Stefan Zohren\n",
    "Institute: Oxford-Man Institute of Quantitative Finance, Department of Engineering Science, University of Oxford\n",
    "Copyright (c) 2021 Oxford Man Institute & University of Oxford. All rights reserved.\n",
    "\n",
    "\n",
    "Copyright (c) 2023 Graphcore Ltd. All rights reserved.\n",
    "This Jupyter notebook has been modified by Graphcore Ltd so that it can be run with the latest version of Graphcore's [Poplar (TM) SDK](https://docs.graphcore.ai/projects/sdk-overview/).\n",
    "\n",
    "On Paperspace:\n",
    "\n",
    "|  Domain | Tasks | Model | Datasets | Workflow |   Number of IPUs   | Execution time |\n",
    "|---------|-------|-------|----------|----------|--------------|--------------|\n",
    "| Finance | prediction  | DeepLOB-Attention |  Limit-Order Books (FI-2010) | training, inference | 4 | ~6 minutes  |\n",
    "\n",
    "**This notebook demonstrates how to train the DeepLOB-Attention model on Graphcore IPUs using TensorFlow 2.**\n",
    "\n",
    "[![Join our Slack Community](https://img.shields.io/badge/Slack-Join%20Graphcore's%20Community-blue?style=flat-square&logo=slack)](https://www.graphcore.ai/join-community)\n",
    "\n",
    "## Methods\n",
    "\n",
    "This Jupyter notebook is used to demonstrate the machine learning methods for multi-horizon forecasting for limit order books, shown in [2], implemented using TensorFlow 2. The publicly available FI-2010 [1] dataset is used for model training, validation and inference.\n",
    "\n",
    "## Data\n",
    "The FI-2010 dataset is publicly available and interested readers can check out the paper [1]. The dataset can be downloaded from: https://etsin.fairdata.fi/dataset/73eb48d7-4dbc-4a10-a52a-da745b47a649 \n",
    "\n",
    "This notebook will download the data automatically. Alternatively, it may be obtained at the following URL: \n",
    "\n",
    "https://drive.google.com/drive/folders/1Xen3aRid9ZZhFqJRgEMyETNazk02cNmv?usp=sharing.\n",
    "\n",
    "## References\n",
    "[1] Ntakaris A, Magris M, Kanniainen J, Gabbouj M, Iosifidis A. Benchmark dataset for mid‐price forecasting of limit order book data with machine learning methods. Journal of Forecasting. 2018 Dec;37(8):852-66. https://arxiv.org/abs/1705.03233\n",
    "\n",
    "[2] Zhang Z, Zohren S. Multi-Horizon Forecasting for Limit Order Books: Novel Deep Learning Approaches and Hardware Acceleration using Intelligent Processing Units. https://arxiv.org/abs/2105.10430\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b58f265b",
   "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-Tensorflow2/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",
    "## Install dependencies\n",
    "\n",
    "First we need to install the required Python libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e442dc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -r requirements.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bc920fa0",
   "metadata": {},
   "source": [
    "## Import the necessary libraries\n",
    "\n",
    "Next we import the Python libraries necessary for running training, validation and inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1fe2cd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import, division, print_function, unicode_literals\n",
    "\n",
    "import os\n",
    "import logging\n",
    "import glob\n",
    "import argparse\n",
    "import sys\n",
    "import time\n",
    "import tensorflow as tf\n",
    "from tensorflow.python import ipu\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import backend as K\n",
    "from ipu_tensorflow_addons.keras.layers import LSTM\n",
    "from sklearn.metrics import accuracy_score, classification_report\n",
    "import pickle\n",
    "import numpy as np\n",
    "from collections import Counter\n",
    "import zipfile"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f9367830",
   "metadata": {},
   "source": [
    "## Download the dataset\n",
    "\n",
    "Now we download the dataset, which should take only a few moments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ded7422",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.isfile(\"data.zip\"):\n",
    "    !wget https://raw.githubusercontent.com/zcakhaa/DeepLOB-Deep-Convolutional-Neural-Networks-for-Limit-Order-Books/master/data/data.zip\n",
    "    print(\"Data downloaded.\")\n",
    "else:\n",
    "    print(\"Data already exists, skipping download.\")\n",
    "\n",
    "with zipfile.ZipFile(\"data.zip\", \"r\") as zip_ref:\n",
    "    zip_ref.extractall(\"./\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "55552f63",
   "metadata": {},
   "source": [
    "## Configuration\n",
    "\n",
    "We set a number of variables relating to model size and hyperparameters, as well as the location for saving the trained model checkpoints. We also set the number of epochs for which we wish to train the model.\n",
    "\n",
    "For full training, it is recommended to set the number of epochs to >= 150 as per the original paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4cb62a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "T = 50  # lookback window size\n",
    "epochs = 20  # number of training epochs\n",
    "batch_size = 16  # gradient descent batch size\n",
    "n_hidden = 64  # hidden state for decoder\n",
    "SHUFFLE = True  # shuffle the traning data\n",
    "saved_model_path = os.getenv(\"CHECKPOINT_DIR\", \"/tmp/checkpoints\")\n",
    "saved_model_path = os.path.join(saved_model_path, \"deeplob_attention\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b6323037",
   "metadata": {},
   "source": [
    "Let's also set an environment variable which allows to use the executable caches, saving us from recompiling the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45cf2019",
   "metadata": {},
   "outputs": [],
   "source": [
    "executable_cache_dir = os.getenv(\"POPLAR_EXECUTABLE_CACHE_DIR\", \"/tmp/exe_cache/\")\n",
    "os.environ[\"TF_POPLAR_FLAGS\"] = f\"--executable_cache_path='{executable_cache_dir}'\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a7ab1ac2",
   "metadata": {},
   "source": [
    "## Data processing\n",
    "\n",
    "Next we define a number of functions which process the data on the host before it is sent to the IPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4c834e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_x(data):\n",
    "    df1 = data[:40, :].T\n",
    "    return np.array(df1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "166b4f98",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_label(data):\n",
    "    lob = data[-5:, :].T\n",
    "    all_label = []\n",
    "\n",
    "    for i in range(lob.shape[1]):\n",
    "        one_label = lob[:, i] - 1\n",
    "        one_label = keras.utils.to_categorical(one_label, 3)\n",
    "        one_label = one_label.reshape(len(one_label), 1, 3)\n",
    "        all_label.append(one_label)\n",
    "\n",
    "    return np.hstack(all_label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2e2e9aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def data_classification(X, Y, T):\n",
    "    [N, D] = X.shape\n",
    "    df = np.array(X)\n",
    "\n",
    "    dY = np.array(Y)\n",
    "\n",
    "    dataY = dY[T - 1 : N]\n",
    "\n",
    "    dataX = np.zeros((N - T + 1, T, D))\n",
    "    for i in range(T, N + 1):\n",
    "        dataX[i - T] = df[i - T : i, :]\n",
    "\n",
    "    return dataX.reshape(dataX.shape + (1,)), dataY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd15f5eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_decoder_input(data, teacher_forcing):\n",
    "    if teacher_forcing:\n",
    "        first_decoder_input = keras.utils.to_categorical(np.zeros(len(data)), 3)\n",
    "        first_decoder_input = first_decoder_input.reshape(\n",
    "            len(first_decoder_input), 1, 3\n",
    "        )\n",
    "        decoder_input_data = np.hstack((data[:, :-1, :], first_decoder_input))\n",
    "\n",
    "    if not teacher_forcing:\n",
    "        decoder_input_data = np.zeros((len(data), 1, 3))\n",
    "        decoder_input_data[:, 0, 0] = 1.0\n",
    "\n",
    "    return decoder_input_data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6f97595a",
   "metadata": {},
   "source": [
    "## Load data from disk\n",
    "\n",
    "Now we load the dataset from the disk into NumPy arrays and pre-process the data using the functions we defined."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b4227a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load data\n",
    "dec_train = np.loadtxt(\"Train_Dst_NoAuction_DecPre_CF_7.txt\")\n",
    "dec_test1 = np.loadtxt(\"Test_Dst_NoAuction_DecPre_CF_7.txt\")\n",
    "dec_test2 = np.loadtxt(\"Test_Dst_NoAuction_DecPre_CF_8.txt\")\n",
    "dec_test3 = np.loadtxt(\"Test_Dst_NoAuction_DecPre_CF_9.txt\")\n",
    "dec_test = np.hstack((dec_test1, dec_test2, dec_test3))\n",
    "\n",
    "# extract limit order book data from the FI-2010 dataset\n",
    "train_lob = prepare_x(dec_train)\n",
    "test_lob = prepare_x(dec_test)\n",
    "\n",
    "# extract label from the FI-2010 dataset\n",
    "train_label = get_label(dec_train)\n",
    "test_label = get_label(dec_test)\n",
    "\n",
    "# prepare training data. We feed past T observations into our algorithms.\n",
    "train_encoder_input, train_decoder_target = data_classification(\n",
    "    train_lob, train_label, T\n",
    ")\n",
    "train_decoder_input = prepare_decoder_input(train_encoder_input, teacher_forcing=False)\n",
    "\n",
    "test_encoder_input, test_decoder_target = data_classification(test_lob, test_label, T)\n",
    "test_decoder_input = prepare_decoder_input(test_encoder_input, teacher_forcing=False)\n",
    "\n",
    "print(\n",
    "    f\"train_encoder_input.shape = {train_encoder_input.shape},\"\n",
    "    f\"train_decoder_target.shape = {train_decoder_target.shape}\"\n",
    ")\n",
    "print(\n",
    "    f\"test_encoder_input.shape = {test_encoder_input.shape},\"\n",
    "    f\"test_decoder_target.shape = {test_decoder_target.shape}\"\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "220d7eaf",
   "metadata": {},
   "source": [
    "## IPU configuration\n",
    "\n",
    "In order to use the IPUs, we create an IPU configuration using the `IPUConfig` class.\n",
    "\n",
    "For this model only one IPU is required to perform training, validation and inference. However by requesting multiple IPUs we can increase throughput by executing tasks in a data parallel fashion.\n",
    "\n",
    "The model is replicated across the IPUs, with an identical copy of the model residing on each IPU. Each IPU receives a different batch of input data, and during training the gradients are automatically averaged across the replicas before the backward pass is performed.\n",
    "\n",
    "By setting `ipu_config.auto_select_ipus = num_ipus` we can automatically select the first available device containing the desired number of IPUs on the system:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1b020c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of IPUs over which to replicate the model\n",
    "available_ipus = int(os.getenv(\"NUM_AVAILABLE_IPU\", 4))\n",
    "num_ipus = min(available_ipus, 4)  # Not intended to scale beyond POD4\n",
    "\n",
    "# Configure the IPU system\n",
    "ipu_config = ipu.config.IPUConfig()\n",
    "ipu_config.auto_select_ipus = num_ipus\n",
    "ipu_config.configure_ipu_system()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "abdede69",
   "metadata": {},
   "source": [
    "## Create the model\n",
    "\n",
    "Now we create the DeepLOB-Attention model, from [Luong et al](https://arxiv.org/abs/1508.04025), using the Keras functional API.\n",
    "\n",
    "The attention model is an evolution of the seq2seq model shown in the DeepLOBSeq2Seq notebook, designed to better handle inputs containing long sequences.\n",
    "\n",
    "Multiple `Input` layers are present since the encoder and decoder parts of the model take separate input tensors, and so when constructing the dataset, this will need to be considered."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaca27b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model_attention(latent_dim):\n",
    "    # Luong Attention\n",
    "    # https://arxiv.org/abs/1508.04025\n",
    "\n",
    "    input_train = keras.Input(shape=(50, 40, 1))\n",
    "\n",
    "    conv_first1 = keras.layers.Conv2D(32, (1, 2), strides=(1, 2))(input_train)\n",
    "    conv_first1 = keras.layers.LeakyReLU(alpha=0.01)(conv_first1)\n",
    "    conv_first1 = keras.layers.Conv2D(32, (4, 1), padding=\"same\")(conv_first1)\n",
    "    conv_first1 = keras.layers.LeakyReLU(alpha=0.01)(conv_first1)\n",
    "    conv_first1 = keras.layers.Conv2D(32, (4, 1), padding=\"same\")(conv_first1)\n",
    "    conv_first1 = keras.layers.LeakyReLU(alpha=0.01)(conv_first1)\n",
    "\n",
    "    conv_first1 = keras.layers.Conv2D(32, (1, 2), strides=(1, 2))(conv_first1)\n",
    "    conv_first1 = keras.layers.LeakyReLU(alpha=0.01)(conv_first1)\n",
    "    conv_first1 = keras.layers.Conv2D(32, (4, 1), padding=\"same\")(conv_first1)\n",
    "    conv_first1 = keras.layers.LeakyReLU(alpha=0.01)(conv_first1)\n",
    "    conv_first1 = keras.layers.Conv2D(32, (4, 1), padding=\"same\")(conv_first1)\n",
    "    conv_first1 = keras.layers.LeakyReLU(alpha=0.01)(conv_first1)\n",
    "\n",
    "    conv_first1 = keras.layers.Conv2D(32, (1, 10))(conv_first1)\n",
    "    conv_first1 = keras.layers.LeakyReLU(alpha=0.01)(conv_first1)\n",
    "    conv_first1 = keras.layers.Conv2D(32, (4, 1), padding=\"same\")(conv_first1)\n",
    "    conv_first1 = keras.layers.LeakyReLU(alpha=0.01)(conv_first1)\n",
    "    conv_first1 = keras.layers.Conv2D(32, (4, 1), padding=\"same\")(conv_first1)\n",
    "    conv_first1 = keras.layers.LeakyReLU(alpha=0.01)(conv_first1)\n",
    "\n",
    "    # build the inception module\n",
    "    convsecond_1 = keras.layers.Conv2D(64, (1, 1), padding=\"same\")(conv_first1)\n",
    "    convsecond_1 = keras.layers.LeakyReLU(alpha=0.01)(convsecond_1)\n",
    "    convsecond_1 = keras.layers.Conv2D(64, (3, 1), padding=\"same\")(convsecond_1)\n",
    "    convsecond_1 = keras.layers.LeakyReLU(alpha=0.01)(convsecond_1)\n",
    "\n",
    "    convsecond_2 = keras.layers.Conv2D(64, (1, 1), padding=\"same\")(conv_first1)\n",
    "    convsecond_2 = keras.layers.LeakyReLU(alpha=0.01)(convsecond_2)\n",
    "    convsecond_2 = keras.layers.Conv2D(64, (5, 1), padding=\"same\")(convsecond_2)\n",
    "    convsecond_2 = keras.layers.LeakyReLU(alpha=0.01)(convsecond_2)\n",
    "\n",
    "    convsecond_3 = keras.layers.MaxPooling2D((3, 1), strides=(1, 1), padding=\"same\")(\n",
    "        conv_first1\n",
    "    )\n",
    "    convsecond_3 = keras.layers.Conv2D(64, (1, 1), padding=\"same\")(convsecond_3)\n",
    "    convsecond_3 = keras.layers.LeakyReLU(alpha=0.01)(convsecond_3)\n",
    "\n",
    "    convsecond_output = keras.layers.concatenate(\n",
    "        [convsecond_1, convsecond_2, convsecond_3], axis=3\n",
    "    )\n",
    "    conv_reshape = keras.layers.Reshape(\n",
    "        (int(convsecond_output.shape[1]), int(convsecond_output.shape[3]))\n",
    "    )(convsecond_output)\n",
    "\n",
    "    # seq2seq\n",
    "    encoder_inputs = conv_reshape\n",
    "    encoder = LSTM(latent_dim, return_state=True, return_sequences=True)\n",
    "    encoder_outputs, state_h, state_c = encoder(encoder_inputs)\n",
    "    states = [state_h, state_c]\n",
    "\n",
    "    # Set up the decoder, which will only process one timestep at a time.\n",
    "    decoder_inputs = keras.Input(shape=(1, 3))\n",
    "    decoder_lstm = LSTM(latent_dim, return_sequences=True, return_state=True)\n",
    "    decoder_dense = keras.layers.Dense(3, activation=\"softmax\", name=\"output_layer\")\n",
    "\n",
    "    all_outputs = []\n",
    "    all_attention = []\n",
    "\n",
    "    encoder_state_h = keras.layers.Reshape((1, int(state_h.shape[1])))(state_h)\n",
    "    inputs = keras.layers.concatenate([decoder_inputs, encoder_state_h], axis=2)\n",
    "\n",
    "    for _ in range(5):\n",
    "        # h'_t = f(h'_{t-1}, y_{t-1}, c)\n",
    "        outputs, state_h, state_c = decoder_lstm(inputs, initial_state=states)\n",
    "        # dot\n",
    "        attention = keras.layers.dot([outputs, encoder_outputs], axes=2)\n",
    "        attention = keras.layers.Activation(\"softmax\")(attention)\n",
    "        # context vector\n",
    "        context = keras.layers.dot([attention, encoder_outputs], axes=[2, 1])\n",
    "        context = keras.layers.BatchNormalization(momentum=0.6)(context)\n",
    "\n",
    "        # y = g(h'_t, c_t)\n",
    "        decoder_combined_context = keras.layers.concatenate([context, outputs])\n",
    "        outputs = decoder_dense(decoder_combined_context)\n",
    "        all_outputs.append(outputs)\n",
    "        all_attention.append(attention)\n",
    "\n",
    "        inputs = keras.layers.concatenate([outputs, context], axis=2)\n",
    "        states = [state_h, state_c]\n",
    "\n",
    "    # Concatenate all predictions\n",
    "    decoder_outputs = keras.layers.Lambda(\n",
    "        lambda x: K.concatenate(x, axis=1), name=\"outputs\"\n",
    "    )(all_outputs)\n",
    "    decoder_attention = keras.layers.Lambda(\n",
    "        lambda x: K.concatenate(x, axis=1), name=\"attentions\"\n",
    "    )(all_attention)\n",
    "\n",
    "    # Define and compile model as previously\n",
    "    model = keras.Model(inputs=[input_train, decoder_inputs], outputs=decoder_outputs)\n",
    "    return model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6a234d5b",
   "metadata": {},
   "source": [
    "## Dataset creation\n",
    "\n",
    "Having constructed the Keras model, the next step is to define the construction of the datasets for training, validation and inference.\n",
    "\n",
    "We define a dataset creation function that produces a TensorFlow `Dataset` object from the input tensors, casts to the `float32` datatype and then batches and optionally shuffles/repeats the dataset depending on whether it is intended for model training, validation or inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "098b8d96",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dataset(\n",
    "    encoder_input, decoder_input, encoder_target, batch_size, method, shuffle=False\n",
    "):\n",
    "    train_pairs_ds = tf.data.Dataset.from_tensor_slices((encoder_input, decoder_input))\n",
    "    train_pairs_ds = train_pairs_ds.map(\n",
    "        lambda d, l: (tf.cast(d, tf.float32), tf.cast(l, tf.float32))\n",
    "    )\n",
    "\n",
    "    train_y_ds = tf.data.Dataset.from_tensor_slices(encoder_target)\n",
    "    train_y_ds = train_y_ds.map(lambda d: (tf.cast(d, tf.float32)))\n",
    "\n",
    "    if method != \"prediction\":\n",
    "        train_ds = tf.data.Dataset.zip((train_pairs_ds, train_y_ds))\n",
    "\n",
    "        if shuffle:\n",
    "            train_ds = train_ds.shuffle(len(encoder_input))\n",
    "        train_ds = train_ds.batch(batch_size, drop_remainder=True)\n",
    "\n",
    "    if method == \"train\":\n",
    "        return train_ds.repeat()\n",
    "\n",
    "    if method == \"val\":\n",
    "        return train_ds\n",
    "\n",
    "    if method == \"prediction\":\n",
    "        test_ds = tf.data.Dataset.from_tensor_slices((encoder_input, decoder_input))\n",
    "        test_ds = test_ds.batch(batch_size, drop_remainder=True)\n",
    "        test_ds = test_ds.map(\n",
    "            lambda d, l: [(tf.cast(d, tf.float32), tf.cast(l, tf.float32))]\n",
    "        )\n",
    "\n",
    "        return test_ds"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b99e4ff1",
   "metadata": {},
   "source": [
    "## Steps per execution\n",
    "\n",
    "The IPU can perform multiple training, validation or inference steps in an on-device loop for a single call to the model. Since this reduces the amount of host-IPU communication, this can allow for a much better throughput for a given process.\n",
    "\n",
    "The number of steps executed in such a loop per call to the underlying hardware is controlled by the `steps_per_execution` variable.\n",
    "\n",
    "Setting this value greater than 1 typically improves performance, whilst having no effect on how the model trains in terms of weight updates. However, we need to ensure the number of batches in the dataset is divisible by this number.\n",
    "\n",
    "We therefore define a helper function to calculate the maximum permissible values for `steps_per_epoch`, `validation_steps` and `test_steps` such that the dataset has enough data for each epoch of training and to run model validation and inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6dead1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "steps_per_execution = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80cda859",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_dataset_divisible(num_elements, batch_size, steps_per_exe, num_replicas):\n",
    "    return (\n",
    "        num_elements\n",
    "        // batch_size\n",
    "        // steps_per_exe\n",
    "        // num_replicas\n",
    "        * num_replicas\n",
    "        * steps_per_exe\n",
    "    )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f503cd24",
   "metadata": {},
   "source": [
    "## Train the model\n",
    "\n",
    "We can now train our model using the Keras `model.fit()` API.\n",
    "\n",
    "An `IPUStrategy` class is created, with model construction and execution done from within the scope of this strategy in order to target the IPU.\n",
    "\n",
    "The `steps_per_execution` argument is added to `model.compile()`, along with the loss, optimiser and metrics.\n",
    "\n",
    "We train the model nominally in chunks of 5 epochs. At the end of each set of 5 epochs, we perform model validation to ensure the model is not overfitting to the training dataset.\n",
    "\n",
    "Once the model has trained for the desired number of epochs, the model checkpoint is saved to the directory we specified earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8025ae62",
   "metadata": {},
   "outputs": [],
   "source": [
    "strategy = ipu.ipu_strategy.IPUStrategy()\n",
    "all_results = [[1000, 0]]\n",
    "split_train_val = int(np.floor(len(train_encoder_input) * 0.8))\n",
    "\n",
    "with strategy.scope():\n",
    "    # Create an instance of the model\n",
    "    model = get_model_attention(n_hidden)\n",
    "\n",
    "    # Get the dataset\n",
    "    train_ds = create_dataset(\n",
    "        train_encoder_input[:split_train_val],\n",
    "        train_decoder_input[:split_train_val],\n",
    "        train_decoder_target[:split_train_val],\n",
    "        batch_size,\n",
    "        method=\"train\",\n",
    "        shuffle=SHUFFLE,\n",
    "    )\n",
    "    val_ds = create_dataset(\n",
    "        train_encoder_input[split_train_val:],\n",
    "        train_decoder_input[split_train_val:],\n",
    "        train_decoder_target[split_train_val:],\n",
    "        batch_size,\n",
    "        method=\"val\",\n",
    "    )\n",
    "    test_ds = create_dataset(\n",
    "        test_encoder_input,\n",
    "        test_decoder_input,\n",
    "        test_decoder_target,\n",
    "        batch_size,\n",
    "        method=\"prediction\",\n",
    "    )\n",
    "\n",
    "    # Train the model\n",
    "    adam = keras.optimizers.Adam(learning_rate=0.00004, beta_1=0.9, beta_2=0.999)\n",
    "\n",
    "    model.compile(\n",
    "        loss=\"categorical_crossentropy\",\n",
    "        metrics=[\"accuracy\"],\n",
    "        optimizer=adam,\n",
    "        steps_per_execution=steps_per_execution,\n",
    "    )\n",
    "\n",
    "    epoch_ = 0\n",
    "    epochs_per_fit = 5\n",
    "\n",
    "    train_batches = len(train_encoder_input[:split_train_val])\n",
    "    val_batches = len(train_encoder_input[split_train_val:])\n",
    "    test_batches = len(test_encoder_input)\n",
    "\n",
    "    steps_per_epoch = make_dataset_divisible(\n",
    "        train_batches, batch_size, steps_per_execution, num_ipus\n",
    "    )\n",
    "    val_steps = make_dataset_divisible(\n",
    "        val_batches, batch_size, steps_per_execution, num_ipus\n",
    "    )\n",
    "    test_steps = make_dataset_divisible(\n",
    "        test_batches, batch_size, steps_per_execution, num_ipus\n",
    "    )\n",
    "\n",
    "    while epoch_ < epochs:\n",
    "\n",
    "        model.fit(\n",
    "            train_ds,\n",
    "            steps_per_epoch=steps_per_epoch,\n",
    "            initial_epoch=epoch_,\n",
    "            epochs=epoch_ + epochs_per_fit,\n",
    "        )\n",
    "\n",
    "        epoch_ = epoch_ + epochs_per_fit\n",
    "        result = model.evaluate(\n",
    "            val_ds,\n",
    "            steps=val_steps,\n",
    "        )\n",
    "\n",
    "        all_results.append(result)\n",
    "        print(f\"Epoch = {epoch_},\" f\"Validation Results = {result}\")\n",
    "\n",
    "        if all_results[-1][0] < all_results[-2][0]:\n",
    "            model.save_weights(saved_model_path)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4908092b",
   "metadata": {},
   "source": [
    "## Model inference\n",
    "\n",
    "Having trained the model, or simply reloaded the model weights from a checkpoint generated previously, we can now run inference over the test dataset.\n",
    "\n",
    "Again, we ensure the model is called from within `strategy.scope()` in order to target the IPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50e09709",
   "metadata": {},
   "outputs": [],
   "source": [
    "with strategy.scope():\n",
    "    model.load_weights(saved_model_path)\n",
    "    pred = model.predict(test_ds, steps=test_steps)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "13daf9d8",
   "metadata": {},
   "source": [
    "## Results\n",
    "\n",
    "Finally we can compare the output obtained from inference with the ground truth values, to see how the model performs. We define a helper function which calculates and prints the accuracy score using the SciKit-Learn library, as well as other metrics of interest, for each prediction horizon."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d7ab627",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluation_metrics(real_y, pred_y):\n",
    "    real_y = real_y[: len(pred_y)]\n",
    "    logging.info(\"-------------------------------\")\n",
    "\n",
    "    for i in range(real_y.shape[1]):\n",
    "        print(f\"Prediction horizon = {i}\")\n",
    "        print(\n",
    "            f\"accuracy_score = {accuracy_score(np.argmax(real_y[:, i], axis=1), np.argmax(pred_y[:, i], axis=1))}\"\n",
    "        )\n",
    "        print(\n",
    "            f\"classification_report = {classification_report(np.argmax(real_y[:, i], axis=1), np.argmax(pred_y[:, i], axis=1), digits=4)}\"\n",
    "        )\n",
    "        print(\"-------------------------------\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac83aa7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluation_metrics(test_decoder_target, pred)\n",
    "ipu.config.reset_ipu_configuration()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7b182d6b",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "We have demonstrated how you can use IPUs to run training and inference on the DeepLOB-Attention model with an accuracy of 80%.\n",
    "\n",
    "Interested in other applications for TensorFlow 2 on the IPU? Check out our GNN notebooks:\n",
    "\n",
    "* GPS++ model found in `/ogb-competition`\n",
    "* Cluster GCN model found in `/gnn-cluster-gcn`\n",
    "\n",
    "We also have tutorials dedicated to using IPUs with TensorFlow 2 which is located in `/learning-Tensorflow2-on-IPU` which includes an MNIST tutorial."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
