ANP_1D_regression.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Demo of 1D regression with an Attentive Neural Process (ANP) model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook will provide a simple and straightforward demonstration on how to utilize an Attentive Neural Process (ANP) to regress context and target points to a sine curve.\n",
    "\n",
    "First, we need to import all necessary packages and modules for our task:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import torch\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "# Provide access to modules in repo.\n",
    "sys.path.insert(0, os.path.abspath('neural_process_models'))\n",
    "sys.path.insert(0, os.path.abspath('misc'))\n",
    "\n",
    "from neural_process_models.anp import ANP_Model\n",
    "from misc.test_sin_regression.Sin_Wave_Data import sin_wave_data, plot_functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `sin_wave_data` class, defined in `misc/test_sin_regression/Sin_Data_Wave.py`, represents the curve that we will try to regress to. From instances of this class, we are able to sample context and target points from the curve to serve as inputs for our neural process.\n",
    "\n",
    "The default parameters of this class will produce a \"ground truth\" curve defined as the sum of the following:\n",
    "1. A sine curve with amplitude 1, frequency 1, and phase 1.\n",
    "2. A sine curve with amplitude 2, frequency 2, and phase 1.\n",
    "3. A measured amount of noise (0.1).\n",
    "\n",
    "Let us create an instance of this class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = sin_wave_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we need to instantiate our model. The ANP model is implemented under the `NeuralProcessModel` class under the file `neural_process_models/attentive_neural_process.py`.\n",
    "\n",
    "We will use the following parameters for our example model:\n",
    "* 1 for x-dimension and y-dimension (since this is 1D regression)\n",
    "* 4 hidden layers of dimension 256 for encoders and decoder\n",
    "* 256 as the latent dimension for encoders and decoder\n",
    "* We will utilize dot attention for the self-attention process.\n",
    "* We will utilize multihead attention for the cross-attention process.\n",
    "* We will utilize a deterministic path for the encoder.\n",
    "\n",
    "Let us create an instance of this class, as well as set some hyperparameters for our training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np_model = ANP_Model(x_dim=1,\n",
    "                     y_dim=1,\n",
    "                     mlp_hidden_size_list=[256, 256, 256, 256],\n",
    "                     latent_dim=256,\n",
    "                     use_rnn=False,\n",
    "                     use_self_attention=False,\n",
    "                     use_deter_path=True)\n",
    "\n",
    "optim = torch.optim.Adam(np_model.parameters(), lr=1e-4)\n",
    "num_epochs = 1000\n",
    "batch_size = 16"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let us train our model. For each epoch, we will print the loss at that epoch. Here is a GIF displaying the regression results over the duration of training:\n",
    "\n",
    "![SegmentLocal](misc/images/anp_-3_3.gif \"results\")\n",
    "\n",
    "Additionally, every 50 epochs, an image will be generated and displayed, using `pyplot`. This will give you an opportunity to more closely analyze and/or save the images, if you would like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Training...\")\n",
    "\n",
    "for epoch in range(1, num_epochs + 1):\n",
    "    print(\"step = \" + str(epoch))\n",
    "\n",
    "    np_model.train()\n",
    "\n",
    "    plt.clf()\n",
    "    optim.zero_grad()\n",
    "\n",
    "    ctt_x, ctt_y, tgt_x, tgt_y = data.query(batch_size=batch_size,\n",
    "                                            context_x_start=-6,\n",
    "                                            context_x_end=6,\n",
    "                                            context_x_num=200,\n",
    "                                            target_x_start=-6,\n",
    "                                            target_x_end=6,\n",
    "                                            target_x_num=200)\n",
    "    \n",
    "    mu, sigma, log_p, kl, loss = np_model(ctt_x, ctt_y, tgt_x, tgt_y)\n",
    "\n",
    "    print('loss = ', loss)\n",
    "    loss.backward()\n",
    "    optim.step()\n",
    "    np_model.eval()\n",
    "    \n",
    "    if epoch % 50 == 0:\n",
    "        plt.ion()\n",
    "        plot_functions(tgt_x.numpy(),\n",
    "                       tgt_y.numpy(),\n",
    "                       ctt_x.numpy(),\n",
    "                       ctt_y.numpy(),\n",
    "                       mu.detach().numpy(),\n",
    "                       sigma.detach().numpy())\n",
    "        title_str = 'Training at epoch ' + str(epoch)\n",
    "        plt.title(title_str)\n",
    "        plt.pause(0.1)\n",
    "\n",
    "plt.ioff()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}