ANPRNN_1D_regression.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Demo of 1D regression with an Attentive Neural Process with Recurrent Neural Network (ANP-RNN) model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook will provide a simple and straightforward demonstration on how to utilize an Attentive Neural Process with a Recurrent Neural Network (ANP-RNN) 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_rnn import ANP_RNN_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 a self-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_RNN_Model(x_dim=1,\n",
    "                     y_dim=1,\n",
    "                     mlp_hidden_size_list=[256, 256, 256, 256],\n",
    "                     latent_dim=256,\n",
    "                     use_rnn=True,\n",
    "                     use_self_attention=True,\n",
    "                     le_self_attention_type=\"laplace\",\n",
    "                     de_self_attention_type=\"laplace\",\n",
    "                     de_cross_attention_type=\"laplace\",\n",
    "                     use_deter_path=True)\n",
    "\n",
    "optim = torch.optim.Adam(np_model.parameters(), lr=1e-4)\n",
    "\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.\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": [
    "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() "
   ]
  }
 ],
 "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": 2
}