ANP_2D_regression.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Demo of 2D 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 perform regression on images in the MNIST dataset.\n",
    "\n",
    "**Note:**: the training time for this model is very lengthy; a GPU is recommended to reduce the training time. Comment back in `.cuda` in the code below if you decide to use one.\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",
    "import mnist\n",
    "import random\n",
    "import numpy as np\n",
    "\n",
    "# Provide access to modules in repo.\n",
    "sys.path.insert(0, os.path.abspath('neural_process_models'))\n",
    "\n",
    "from neural_process_models.anp import ANP_Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each data point (image) in the MNIST dataset is represented by a 2D (28 x 28) array, with each cell value existing in the range 0-255.\n",
    "\n",
    "Let us retrieve the prepare this dataset for training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_images = mnist.test_images()  # (10000 x 28 x 28)\n",
    "test_images = (test_images / 255.0)  # normalize pixel values\n",
    "\n",
    "data_size = len(test_images)\n",
    "test_images = np.resize(test_images, (10000, 28, 28, 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that we normalized the pixel values above. The normalized pixel values will be the y-values of the context points in the training dataset, while x-values represent a pair of pixel indices indicating where the pixel with the corresponding y-value is found.\n",
    "\n",
    "Let us initialize 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",
    "* 2 for x-dimension and 1 for y-dimension (as explained above)\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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ANP_Model(x_dim=2,  # x_dim: pixel index (0-27 x 0-27)\n",
    "                  y_dim=1,  # y_dim: normalized pixel value (0-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)#.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And let's set some hyperparameters for our tuning:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optim = torch.optim.Adam(model.parameters(), lr=1e-4)\n",
    "\n",
    "num_epochs = 10000\n",
    "batch_size = 16\n",
    "num_context = 400"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let us train our model. For each epoch, we will print the loss at that epoch."
   ]
  },
  {
   "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",
    "    model.train()\n",
    "\n",
    "    plt.clf()\n",
    "    optim.zero_grad()\n",
    "\n",
    "    ctt_x, ctt_y, tgt_x, tgt_y = list(), list(), list(), list()\n",
    "\n",
    "    sample_context_indices = random.sample(range(data_size), batch_size)\n",
    "\n",
    "    for context_idx in sample_context_indices:\n",
    "        pixel_indices = random.sample(range(784), num_context)\n",
    "\n",
    "        c_x, c_y = list(), list()\n",
    "        for pixel_idx in pixel_indices:\n",
    "            pixel_x = (pixel_idx // 28) / 27.0\n",
    "            pixel_y = (pixel_idx % 28) / 27.0\n",
    "\n",
    "    \t\tc_x.append([pixel_x, pixel_y])\n",
    "    \t\tc_y.append(test_images[context_idx][pixel_x][pixel_y])\n",
    "\n",
    "    \tctt_x.append(c_x)\n",
    "    \tctt_y.append(c_y)\n",
    "\n",
    "    sample_target_indices = random.sample(range(data_size), batch_size)\n",
    "\n",
    "    for target_idx in sample_target_indices:\n",
    "        t_x, t_y = list(), list()\n",
    "        for pixel_x in range(28):\n",
    "        \tfor pixel_y in range(28):\n",
    "\t            t_x.append([pixel_x, pixel_y])\n",
    "\t            t_y.append(test_images[target_idx][pixel_x][pixel_y])\n",
    "\n",
    "        tgt_x.append(t_x)\n",
    "        tgt_y.append(t_y)\n",
    "\n",
    "    ctt_x = torch.FloatTensor(ctt_x)#.cuda()\n",
    "    ctt_y = torch.FloatTensor(ctt_y)#.cuda()\n",
    "    tgt_x = torch.FloatTensor(tgt_x)#.cuda()\n",
    "    tgt_y = torch.FloatTensor(tgt_y)#.cuda()\n",
    "\n",
    "\n",
    "    # ctt_x: (batch_size x num_context x 2), ctt_y: (batch_size x 784 x 1)\n",
    "    # tgt_x: (batch_size x num_context x 2), tgt_y: (batch_size x 784 x 1)\n",
    "    mu, sigma, log_p, kl, loss = model(ctt_x, ctt_y, tgt_x, tgt_y)\n",
    "\n",
    "    # print('kl =', kl)\n",
    "    print('loss = ', loss)\n",
    "    # print('mu.size() =', mu.size())\n",
    "    # print('sigma.size() =', sigma.size())\n",
    "\n",
    "    # tgt_x_np = tgt_x[0, :, :].squeeze(-1).numpy()\n",
    "    # print('tgt_x_np.shape =', tgt_x_np.shape)\n",
    "\n",
    "    loss.backward()\n",
    "    optim.step()\n",
    "\n",
    "    model.eval()\n",
    "    plt.ion()\n",
    "    # fig = plt.figure()\n",
    "    \n",
    "    # Visualize first target image.\n",
    "    pred_y = mu[0].view(28, 28).detach().numpy()\n",
    "\n",
    "    plt.axis('off')\n",
    "    #plt.imshow(torch.sigmoid(tgt_y).squeeze(0).view(-1, 28).detach().numpy())\n",
    "    plt.imshow(pred_y)\n",
    "\n",
    "    title_str = 'Training at epoch ' + str(epoch)\n",
    "    plt.title(title_str)\n",
    "    plt.savefig(title_str + \".png\")\n",
    "    plt.pause(0.1)\n",
    "\n",
    "plt.ioff()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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": 2
}