06_04_EXP_TransfLearn_MissingMNIST_ConvTimeLSTM.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import torch\n",
    "from torch.autograd import Variable\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import moving MNIST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Moving_MNIST = np.load('data/mnist_test_seq.npy')\n",
    "Moving_MNIST = Moving_MNIST / 255\n",
    "Moving_MNIST.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Give `torch` the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Making into PyTorch tensor\n",
    "Moving_MNIST_tensor = torch.from_numpy(Moving_MNIST)\n",
    "\n",
    "# Putting the existing dimensions into appropriate order\n",
    "Moving_MNIST_tensor = Moving_MNIST_tensor.permute(1, 0, 2, 3)\n",
    "\n",
    "# Added the acknowledge that this is 1 spectral band\n",
    "Moving_MNIST_tensor = Moving_MNIST_tensor.unsqueeze(2)\n",
    "\n",
    "# Checking shape\n",
    "Moving_MNIST_tensor.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train/validation split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_indices = np.random.choice(range(10000), size = 8000, replace = False)\n",
    "\n",
    "OutofSample_indices = [index for index in range(10000) if index not in train_indices.tolist()]\n",
    "validation_indices = np.random.choice(OutofSample_indices, size = 1000, replace = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating 9 data subsets, which the data will be uniformly sampled into and subsetted from 20 frames to [11, 20] frames, keep tracking of the number of frames omitted (i.e. $\\Delta{t}$ where 1 = no omission, 2 = 1 omission, etc...)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating empty nested lists to store data...\n",
    "train_data = []\n",
    "for i in range(9):\n",
    "    train_data.append([])\n",
    "# ...and time steps in\n",
    "train_delta_Ts = []\n",
    "for i in range(9):\n",
    "    train_delta_Ts.append([])\n",
    "\n",
    "# For the train (above) and validation (below) data\n",
    "validation_data = []\n",
    "for i in range(9):\n",
    "    validation_data.append([])\n",
    "validation_delta_Ts = []\n",
    "for i in range(9):\n",
    "    validation_delta_Ts.append([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(Moving_MNIST_tensor.shape[0]):\n",
    "    \n",
    "    # Determine how many frames to keep out of the 20 possible\n",
    "    #     leaving 11 so that there is at least 10 for input and 1 for out\n",
    "    num_frames = np.random.choice(range(11, 20))\n",
    "    \n",
    "    # Determining what index that places the data in\n",
    "    data_index = (num_frames - 11)\n",
    "    \n",
    "    # Now that we have a number to keep, picking which specific ones to keep\n",
    "    frame_indices = np.random.choice(range(20),\n",
    "                                     size = num_frames,\n",
    "                                     replace = False)\n",
    "    \n",
    "    # Sorting those frames to proper chronological order\n",
    "    sorted_frame_indices = sorted(frame_indices)\n",
    "    \n",
    "    # Getting the missing-frames-data\n",
    "    frame_data = Moving_MNIST_tensor[i, sorted_frame_indices]\n",
    "    frame_data = frame_data.unsqueeze(dim = 0)\n",
    "    \n",
    "    # Determining \"time steps\" between indices\n",
    "    #     1 = proper sequence\n",
    "    #     2 = missing one index inbetween\n",
    "    #     3 = missing two indices inbetween\n",
    "    #     etc...\n",
    "    delta_Ts = []\n",
    "    for j in range(1, len(sorted_frame_indices)):\n",
    "        delta_Ts.append(sorted_frame_indices[j] - sorted_frame_indices[j-1])\n",
    "        \n",
    "    # Getting delta_Ts as an image band\n",
    "    #     \"num_frames - 1\" because it's not necessary for the first band\n",
    "    #     since it has no prior reference\n",
    "    delta_T_tensors = torch.ones([1, num_frames - 1, 1, 64, 64])\n",
    "    for k in range(delta_T_tensors.shape[1]):\n",
    "        delta_T_tensors[:, k, :, :, :] *= delta_Ts[k]\n",
    "    \n",
    "    # Storing it and its time steps appropriately\n",
    "    if i in train_indices:\n",
    "        # Initiating the tensor within the empty list\n",
    "        if train_data[data_index] == []:\n",
    "            train_data[data_index].append(frame_data)\n",
    "            train_delta_Ts[data_index].append(delta_T_tensors)\n",
    "        # Or expanding upon the tensor\n",
    "        else:\n",
    "            train_data[data_index][0] = torch.cat([train_data[data_index][0],\n",
    "                                                   frame_data])\n",
    "            train_delta_Ts[data_index][0] = torch.cat([train_delta_Ts[data_index][0],\n",
    "                                                       delta_T_tensors])\n",
    "    elif i in validation_indices:\n",
    "        if validation_data[data_index] == []:\n",
    "            validation_data[data_index].append(frame_data)\n",
    "            validation_delta_Ts[data_index].append(delta_T_tensors)\n",
    "        else:\n",
    "            validation_data[data_index][0] = torch.cat([validation_data[data_index][0],\n",
    "                                                        frame_data])\n",
    "            validation_delta_Ts[data_index][0] = torch.cat([validation_delta_Ts[data_index][0],\n",
    "                                                            delta_T_tensors])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('For train data...')\n",
    "for i in range(len(train_data)):\n",
    "    print(train_data[i][0].shape)\n",
    "\n",
    "print('\\nFor validation data...')\n",
    "for i in range(len(validation_data)):\n",
    "    print(validation_data[i][0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Separating $x$ and $y$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sep_x_y(tensor_to_separate, time_tensor_to_separate):\n",
    "    x = []\n",
    "    y = []\n",
    "    t = []\n",
    "    for i in range(len(tensor_to_separate)):\n",
    "        current_seq = tensor_to_separate[i, :, :, :, :]\n",
    "        current_time_seq = time_tensor_to_separate[i, :, :, :, :]\n",
    "        for j in range(current_seq.shape[0]):\n",
    "            if j >= 10:\n",
    "                current_x = current_seq[(j-10):j].numpy()\n",
    "                x.append(current_x)\n",
    "                current_y = current_seq[j].unsqueeze(dim = 0).numpy()\n",
    "                y.append(current_y)\n",
    "                current_t = current_time_seq[(j-10):j].numpy()\n",
    "                t.append(current_t)\n",
    "    x = np.asarray(x)\n",
    "    x = torch.from_numpy(x).type(torch.FloatTensor)\n",
    "    y = np.asarray(y)\n",
    "    y = torch.from_numpy(y).type(torch.FloatTensor)\n",
    "    t = np.asarray(t)\n",
    "    t = torch.from_numpy(t).type(torch.FloatTensor)\n",
    "    return(x, y, t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_11, y_11, t_11 = sep_x_y(train_data[0][0], train_delta_Ts[0][0])\n",
    "x_12, y_12, t_12 = sep_x_y(train_data[1][0], train_delta_Ts[1][0])\n",
    "x_13, y_13, t_13 = sep_x_y(train_data[2][0], train_delta_Ts[2][0])\n",
    "x_14, y_14, t_14 = sep_x_y(train_data[3][0], train_delta_Ts[3][0])\n",
    "x_15, y_15, t_15 = sep_x_y(train_data[4][0], train_delta_Ts[4][0])\n",
    "x_16, y_16, t_16 = sep_x_y(train_data[5][0], train_delta_Ts[5][0])\n",
    "x_17, y_17, t_17 = sep_x_y(train_data[6][0], train_delta_Ts[6][0])\n",
    "x_18, y_18, t_18 = sep_x_y(train_data[7][0], train_delta_Ts[7][0])\n",
    "x_19, y_19, t_19 = sep_x_y(train_data[8][0], train_delta_Ts[8][0])\n",
    "\n",
    "x_11_validation, y_11_validation, t_11_validation = sep_x_y(validation_data[0][0],\n",
    "                                                            validation_delta_Ts[0][0])\n",
    "x_12_validation, y_12_validation, t_12_validation = sep_x_y(validation_data[1][0],\n",
    "                                                            validation_delta_Ts[1][0])\n",
    "x_13_validation, y_13_validation, t_13_validation = sep_x_y(validation_data[2][0],\n",
    "                                                            validation_delta_Ts[2][0])\n",
    "x_14_validation, y_14_validation, t_14_validation = sep_x_y(validation_data[3][0],\n",
    "                                                            validation_delta_Ts[3][0])\n",
    "x_15_validation, y_15_validation, t_15_validation = sep_x_y(validation_data[4][0],\n",
    "                                                            validation_delta_Ts[4][0])\n",
    "x_16_validation, y_16_validation, t_16_validation = sep_x_y(validation_data[5][0],\n",
    "                                                            validation_delta_Ts[5][0])\n",
    "x_17_validation, y_17_validation, t_17_validation = sep_x_y(validation_data[6][0],\n",
    "                                                            validation_delta_Ts[6][0])\n",
    "x_18_validation, y_18_validation, t_18_validation = sep_x_y(validation_data[7][0],\n",
    "                                                            validation_delta_Ts[7][0])\n",
    "x_19_validation, y_19_validation, t_19_validation = sep_x_y(validation_data[8][0],\n",
    "                                                            validation_delta_Ts[8][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Consolidating the 9 subsets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.cat([x_11, x_12, x_13,\n",
    "               x_14, x_15, x_16,\n",
    "               x_17, x_18, x_19])\n",
    "t = torch.cat([t_11, t_12, t_13,\n",
    "               t_14, t_15, t_16,\n",
    "               t_17, t_18, t_19])\n",
    "y = torch.cat([y_11, y_12, y_13,\n",
    "               y_14, y_15, y_16,\n",
    "               y_17, y_18, y_19])\n",
    "\n",
    "x_validation = torch.cat([x_11_validation, x_12_validation, x_13_validation,\n",
    "                          x_14_validation, x_15_validation, x_16_validation,\n",
    "                          x_17_validation, x_18_validation, x_19_validation])\n",
    "t_validation = torch.cat([t_11_validation, t_12_validation, t_13_validation,\n",
    "                          t_14_validation, t_15_validation, t_16_validation,\n",
    "                          t_17_validation, t_18_validation, t_19_validation])\n",
    "y_validation = torch.cat([y_11_validation, y_12_validation, y_13_validation,\n",
    "                          y_14_validation, y_15_validation, y_16_validation,\n",
    "                          y_17_validation, y_18_validation, y_19_validation])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Viewing the overlaid $(x, y)$ sequences where blue = early $x$ frame, green = late $x$ frame, and red = $y$ frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_ex():\n",
    "    random_index = np.random.choice(len(x_validation))\n",
    "    for i in range(10):\n",
    "        if i > 5:\n",
    "            plt.imshow(x_validation[random_index, i, 0], alpha = 0.25, cmap = 'Greens')\n",
    "            plt.text(65, 60/10*i, str(t_validation[random_index, i, 0][0, 0].item()) + ' steps')\n",
    "        else:\n",
    "            plt.imshow(x_validation[random_index, i, 0], alpha = 0.5, cmap = 'Blues')\n",
    "            plt.text(65, 60/10*i, str(t_validation[random_index, i, 0][0, 0].item()) + ' steps')\n",
    "    plt.imshow(y_validation[random_index, 0, 0], cmap = 'Reds', alpha = 0.25)\n",
    "    plt.pause(0.01);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_ex()\n",
    "plot_ex()\n",
    "plot_ex()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Defining the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvTime_LSTMCell(nn.Module):\n",
    "\n",
    "    def __init__(self, input_size, input_dim, hidden_dim, kernel_size, bias, GPU):\n",
    "        \"\"\"\n",
    "        Initialize ConvTime_LSTM cell.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        input_size: (int, int)\n",
    "            Height and width of input tensor as (height, width).\n",
    "        input_dim: int\n",
    "            Number of channels of input tensor.\n",
    "        hidden_dim: int\n",
    "            Number of channels of hidden state.\n",
    "        kernel_size: (int, int)\n",
    "            Size of the convolutional kernel.\n",
    "        bias: bool\n",
    "            Whether or not to add the bias.\n",
    "        \"\"\"\n",
    "\n",
    "        super(ConvTime_LSTMCell, self).__init__()\n",
    "\n",
    "        self.height, self.width = input_size\n",
    "        self.input_dim  = input_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "\n",
    "        self.kernel_size = kernel_size\n",
    "        self.padding     = kernel_size[0] // 2, kernel_size[1] // 2\n",
    "        self.bias        = bias\n",
    "        self.GPU         = GPU\n",
    "        \n",
    "        ## Defining the input convolutional layer ##\n",
    "        self.i_conv = nn.Conv2d(in_channels=self.input_dim + self.hidden_dim,\n",
    "                                out_channels=self.hidden_dim,\n",
    "                                kernel_size=self.kernel_size,\n",
    "                                padding=self.padding,\n",
    "                                bias=self.bias)\n",
    "        \n",
    "        ## Defining the T2 convolutional layer ##\n",
    "        self.T1_conv_x = nn.Conv2d(in_channels=self.input_dim,\n",
    "                                   out_channels=self.hidden_dim,\n",
    "                                   kernel_size=self.kernel_size,\n",
    "                                   padding=self.padding,\n",
    "                                   bias=self.bias)\n",
    "        self.T1_conv_t = nn.Conv2d(in_channels=1,\n",
    "                                   out_channels=self.hidden_dim,\n",
    "                                   kernel_size=self.kernel_size,\n",
    "                                   padding=self.padding,\n",
    "                                   bias=self.bias)\n",
    "        \n",
    "        ## Defining the T1 convolutional layer ##\n",
    "        self.T2_conv_x = nn.Conv2d(in_channels=self.input_dim,\n",
    "                                   out_channels=self.hidden_dim,\n",
    "                                   kernel_size=self.kernel_size,\n",
    "                                   padding=self.padding,\n",
    "                                   bias=self.bias)\n",
    "        self.T2_conv_t = nn.Conv2d(in_channels=1,\n",
    "                                   out_channels=self.hidden_dim,\n",
    "                                   kernel_size=self.kernel_size,\n",
    "                                   padding=self.padding,\n",
    "                                   bias=self.bias)\n",
    "        \n",
    "        ## Defining the activation convolutional layer ##\n",
    "        self.c_conv = nn.Conv2d(in_channels=self.input_dim + self.hidden_dim,\n",
    "                                out_channels=self.hidden_dim,\n",
    "                                kernel_size=self.kernel_size,\n",
    "                                padding=self.padding,\n",
    "                                bias=self.bias)\n",
    "        \n",
    "        ## Defining the output convolutional layer ##\n",
    "        self.o_conv = nn.Conv2d(in_channels=self.input_dim + self.hidden_dim + 1,\n",
    "                                out_channels=self.hidden_dim,\n",
    "                                kernel_size=self.kernel_size,\n",
    "                                padding=self.padding,\n",
    "                                bias=self.bias)\n",
    "\n",
    "    def forward(self, input_tensor, time_tensor, cur_state):\n",
    "        \n",
    "        \n",
    "        ## Getting the h_{m-1} and c_{m-1} ##\n",
    "        ##     the previous hidden and activations ##\n",
    "        h_cur, c_cur = cur_state\n",
    "\n",
    "\n",
    "        ## concatenate the prev. hidden state and the current input along the color channel dim ##\n",
    "        x_h_combined = torch.cat([input_tensor, h_cur], dim = 1)\n",
    "        x_h_t_combined = torch.cat([input_tensor, h_cur, time_tensor], dim = 1)\n",
    "        \n",
    "        \n",
    "        ## The input gate ##\n",
    "        ## Running the input convolution ##\n",
    "        i_conv_outputs = self.i_conv(x_h_combined)\n",
    "        ## Running the input LSTM gate equations ##\n",
    "        i_m = torch.sigmoid(i_conv_outputs)\n",
    "        \n",
    "        \n",
    "        ## The first time gate ##\n",
    "        ## Running the first time convolution for x ##\n",
    "        T1_x_conv_output = self.T1_conv_x(input_tensor)\n",
    "        ## Running the first time convolution for t ##\n",
    "        ## Ensuring that the theoretical constraint of non-positive is met ##\n",
    "        self.T1_conv_t.weight = torch.nn.Parameter(self.T1_conv_t.weight.clamp(max = 0))\n",
    "        ## Passing the convolution ##\n",
    "        T1_t_conv_output = self.T1_conv_t(time_tensor)\n",
    "        ## Performing the internally nested non-linearity ##\n",
    "        T1_t_conv_output = torch.tanh(T1_t_conv_output)\n",
    "        ## Consolidating the output of image and time ##\n",
    "        T1_conv_outputs = T1_t_conv_output + T1_x_conv_output\n",
    "        ## Performing the externally nested non-linearity ##\n",
    "        T1_m = torch.sigmoid(T1_conv_outputs)\n",
    "        \n",
    "        \n",
    "        ## The second time gate ##\n",
    "        ## Running the second time convolution for x ##\n",
    "        T2_x_conv_output = self.T2_conv_x(input_tensor)\n",
    "        ## Running the first time convolution for t ##\n",
    "        ## Passing the convolution ##\n",
    "        T2_t_conv_output = self.T2_conv_t(time_tensor)\n",
    "        ## Performing the internally nested non-linearity ##\n",
    "        T2_t_conv_output = torch.tanh(T2_t_conv_output)\n",
    "        ## Consolidating the output of image and time ##\n",
    "        T2_conv_outputs = T2_t_conv_output + T2_x_conv_output\n",
    "        ## Performing the externally nested non-linearity ##\n",
    "        T2_m = torch.sigmoid(T2_conv_outputs)\n",
    "        \n",
    "        \n",
    "        ## The c vectors ##\n",
    "        ## Running the c convolution ##\n",
    "        c_conv_outputs = self.c_conv(x_h_combined)\n",
    "        ## Computing the c tilde and c activation vectors ##\n",
    "        c_m_tilde = (((1 - i_m * T1_m) * c_cur) +\n",
    "                     (i_m * T1_m * torch.tanh(c_conv_outputs)))\n",
    "        c_m = (((1 - i_m) * c_cur) +\n",
    "               (i_m * T2_m * torch.tanh(c_conv_outputs)))\n",
    "        \n",
    "         \n",
    "        ## The output gate ##\n",
    "        ## Running the output gate convolution ##\n",
    "        o_conv_output = self.o_conv(x_h_t_combined)\n",
    "        ## Running the output LSTM gate equations ##\n",
    "        o_m = torch.sigmoid(o_conv_output)\n",
    "        \n",
    "        \n",
    "        ## The hidden vector ##\n",
    "        h_m = o_m * torch.tanh(c_m_tilde)\n",
    "        \n",
    "        \n",
    "        return h_m, c_m\n",
    "\n",
    "    def init_hidden(self, batch_size):\n",
    "        to_return = (Variable(torch.zeros(batch_size, self.hidden_dim, self.height, self.width)),\n",
    "                     Variable(torch.zeros(batch_size, self.hidden_dim, self.height, self.width)))\n",
    "        if self.GPU:\n",
    "            to_return = (to_return[0].cuda(), to_return[1].cuda())\n",
    "        return(to_return)\n",
    "\n",
    "\n",
    "class ConvTime_LSTM(nn.Module):\n",
    "\n",
    "    def __init__(self, input_size, input_dim, hidden_dim, kernel_size, num_layers,\n",
    "                 batch_first, bias, return_all_layers, GPU):\n",
    "        super(ConvTime_LSTM, self).__init__()\n",
    "\n",
    "        self._check_kernel_size_consistency(kernel_size)\n",
    "\n",
    "        # Make sure that both `kernel_size` and `hidden_dim` are lists having len == num_layers\n",
    "        kernel_size = self._extend_for_multilayer(kernel_size, num_layers)\n",
    "        hidden_dim  = self._extend_for_multilayer(hidden_dim, num_layers)\n",
    "        if not len(kernel_size) == len(hidden_dim) == num_layers:\n",
    "            raise ValueError('Inconsistent list length.')\n",
    "\n",
    "        self.height, self.width = input_size\n",
    "\n",
    "        self.input_dim  = input_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.kernel_size = kernel_size\n",
    "        self.num_layers = num_layers\n",
    "        self.batch_first = batch_first\n",
    "        self.bias = bias\n",
    "        self.return_all_layers = return_all_layers\n",
    "        self.GPU = GPU\n",
    "\n",
    "        cell_list = []\n",
    "        for i in range(0, self.num_layers):\n",
    "            cur_input_dim = self.input_dim if i == 0 else self.hidden_dim[i-1]\n",
    "\n",
    "            cell_list.append(ConvTime_LSTMCell(input_size=(self.height, self.width),\n",
    "                                               input_dim=cur_input_dim,\n",
    "                                               hidden_dim=self.hidden_dim[i],\n",
    "                                               kernel_size=self.kernel_size[i],\n",
    "                                               bias=self.bias,\n",
    "                                               GPU=self.GPU))\n",
    "\n",
    "        self.cell_list = nn.ModuleList(cell_list)\n",
    "\n",
    "    def forward(self, input_tensor, time_tensor, hidden_state=None):\n",
    "        \"\"\"\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        input_tensor: todo \n",
    "            5-D Tensor either of shape (t, b, c, h, w) or (b, t, c, h, w)\n",
    "        hidden_state: todo\n",
    "            None. todo implement stateful\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        last_state_list, layer_output\n",
    "        \"\"\"\n",
    "        \n",
    "        if not self.batch_first:\n",
    "            # (t, b, c, h, w) -> (b, t, c, h, w)\n",
    "            input_tensor.permute(1, 0, 2, 3, 4)\n",
    "\n",
    "        # Implement stateful ConvLSTM\n",
    "        if hidden_state is not None:\n",
    "            raise NotImplementedError()\n",
    "        else:\n",
    "            hidden_state = self._init_hidden(batch_size=input_tensor.size(0))\n",
    "\n",
    "        layer_output_list = []\n",
    "        last_state_list   = []\n",
    "\n",
    "        seq_len = input_tensor.size(1)\n",
    "        cur_layer_input = input_tensor\n",
    "        cur_time_input = time_tensor\n",
    "\n",
    "        for layer_idx in range(self.num_layers):\n",
    "            h, c = hidden_state[layer_idx]\n",
    "            output_inner = []\n",
    "            for t in range(seq_len):\n",
    "\n",
    "                h, c = self.cell_list[layer_idx](input_tensor = cur_layer_input[:, t, :, :, :],\n",
    "                                                 time_tensor = cur_time_input[:, t, :, :, :],\n",
    "                                                 cur_state=[h, c])\n",
    "                output_inner.append(h)\n",
    "\n",
    "            layer_output = torch.stack(output_inner, dim=1)\n",
    "            cur_layer_input = layer_output\n",
    "\n",
    "            layer_output_list.append(layer_output)\n",
    "            last_state_list.append([h, c])\n",
    "\n",
    "        if not self.return_all_layers:\n",
    "            layer_output_list = layer_output_list[-1:]\n",
    "            last_state_list   = last_state_list[-1:]\n",
    "\n",
    "        return layer_output_list, last_state_list\n",
    "\n",
    "    def _init_hidden(self, batch_size):\n",
    "        init_states = []\n",
    "        for i in range(self.num_layers):\n",
    "            init_states.append(self.cell_list[i].init_hidden(batch_size))\n",
    "        return init_states\n",
    "\n",
    "    @staticmethod\n",
    "    def _check_kernel_size_consistency(kernel_size):\n",
    "        if not (isinstance(kernel_size, tuple) or\n",
    "                    (isinstance(kernel_size, list) and all([isinstance(elem, tuple) for elem in kernel_size]))):\n",
    "            raise ValueError('`kernel_size` must be tuple or list of tuples')\n",
    "\n",
    "    @staticmethod\n",
    "    def _extend_for_multilayer(param, num_layers):\n",
    "        if not isinstance(param, list):\n",
    "            param = [param] * num_layers\n",
    "        return param"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The model already trained on standard MovingMNIST was saved as a parallel GPU PyTorch model, so it automatically loads back as such"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conv_time_lstm = torch.load('ConvTimeLSTMSavedModel2.pt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss = torch.nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(conv_time_lstm.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils import data\n",
    "\n",
    "class train_Dataset(data.Dataset):\n",
    "    'Characterizes a dataset for PyTorch'\n",
    "    def __init__(self, data_indices):\n",
    "        'Initialization'\n",
    "        self.data_indices = data_indices\n",
    "    \n",
    "    def __len__(self):\n",
    "        'Denotes the total number of samples'\n",
    "        return len(self.data_indices)\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        'Generates one sample of data'\n",
    "        # Select sample\n",
    "        IDs = self.data_indices[index]\n",
    "\n",
    "        # Load data and get label\n",
    "        curr_x = x[IDs, :, :, :, :]\n",
    "        curr_t = t[IDs, :, :, :, :] * 1/10\n",
    "        curr_y = y[IDs, :, :, :, :]\n",
    "\n",
    "        #return X, y\n",
    "        return(curr_x, curr_y, curr_t)\n",
    "    \n",
    "class validation_Dataset(data.Dataset):\n",
    "    'Characterizes a dataset for PyTorch'\n",
    "    def __init__(self, data_indices):\n",
    "        'Initialization'\n",
    "        self.data_indices = data_indices\n",
    "    \n",
    "    def __len__(self):\n",
    "        'Denotes the total number of samples'\n",
    "        return len(self.data_indices)\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        'Generates one sample of data'\n",
    "        # Select sample\n",
    "        IDs = self.data_indices[index]\n",
    "\n",
    "        # Load data and get label\n",
    "        curr_x = x_validation[IDs, :, :, :, :]\n",
    "        curr_t = t_validation[IDs, :, :, :, :] * 1/10\n",
    "        curr_y = y_validation[IDs, :, :, :, :]\n",
    "\n",
    "        #return X, y\n",
    "        return(curr_x, curr_y, curr_t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_set = train_Dataset(data_indices = range(y.shape[0]))\n",
    "validation_set = validation_Dataset(data_indices = range(y_validation.shape[0]))\n",
    "\n",
    "batch_size = 64\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(dataset = training_set,\n",
    "                                           batch_size = batch_size,\n",
    "                                           shuffle = True)\n",
    "validation_loader = torch.utils.data.DataLoader(dataset = validation_set,\n",
    "                                                batch_size = batch_size,\n",
    "                                                shuffle = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_list = []\n",
    "epochs = int(np.ceil((7*10**5) / x.shape[0]))\n",
    "for i in range(epochs):\n",
    "    for data in train_loader:\n",
    "\n",
    "        # data loader\n",
    "        batch_x, batch_y, batch_t = data\n",
    "\n",
    "        # move to GPU\n",
    "        batch_x = batch_x.to(device)\n",
    "        batch_t = batch_t.to(device)\n",
    "        batch_y = batch_y.to(device)\n",
    "\n",
    "        # run model and get the prediction\n",
    "        batch_y_hat = conv_time_lstm(batch_x,\n",
    "                                     batch_t)\n",
    "        batch_y_hat = batch_y_hat[0][0][:, -2:-1, :, :, :]\n",
    "\n",
    "        # calculate and store the loss\n",
    "        batch_loss = loss(batch_y, batch_y_hat)\n",
    "        loss_list.append(batch_loss.item())\n",
    "\n",
    "        # update parameters\n",
    "        optimizer.zero_grad()\n",
    "        batch_loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    print('Epoch: ', i, '\\n\\tBatch loss: ', batch_loss.item(), '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(loss_list,\n",
    "         color = 'cyan',\n",
    "         label = 'Batch')\n",
    "plt.plot(np.convolve(loss_list, 1/10*np.ones(10))[10:-10],\n",
    "         color = 'navy',\n",
    "         label = 'Running average')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting random predictions for the validation set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rand_x, rand_y, rand_t = next(iter(validation_loader))\n",
    "\n",
    "rand_y_hat = conv_time_lstm(rand_x.to(device),\n",
    "                            rand_t.to(device))[0][0][:, -2:-1, :, :, :]\n",
    "rand_y_hat = rand_y_hat.cpu().data.numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# View prediction on sequence ($\\hat{y}$ on $x$)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_random_validation_pred():\n",
    "    f, axarr = plt.subplots(1, 2)\n",
    "    f.set_figheight(4)\n",
    "    f.set_figwidth(8)\n",
    "    random_index = np.random.choice(len(rand_x))\n",
    "    for i in range(10):\n",
    "        axarr[0].imshow(rand_x[random_index, i, 0], alpha = 0.25, cmap = 'gist_gray')\n",
    "        axarr[1].imshow(rand_x[random_index, i, 0], alpha = 0.25, cmap = 'gist_gray')\n",
    "        axarr[1].text(65,\n",
    "                      60/10*i,\n",
    "                      str(int(rand_t[random_index, i, 0, 0, 0].item()*10)) + ' steps')\n",
    "    axarr[0].imshow(rand_y[random_index, 0, 0], cmap = 'Reds', alpha = 0.5)\n",
    "    axarr[0].set_title('Red = True', fontsize = 15)\n",
    "    axarr[1].imshow(rand_y_hat[random_index, 0, 0], cmap = 'Blues', alpha = 0.5)\n",
    "    axarr[1].set_title('Blue = Predicted', fontsize = 15);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_random_validation_pred()\n",
    "plot_random_validation_pred()\n",
    "plot_random_validation_pred()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# View prediction on true ($\\hat{y}$ on $y$)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def overlay_pred_true():\n",
    "    random_index = np.random.choice(len(rand_x))\n",
    "    plt.imshow(rand_y[random_index, 0, 0], cmap = 'Reds', alpha = 0.5)\n",
    "    plt.imshow(rand_y_hat[random_index, 0, 0], cmap = 'Blues', alpha = 0.5)\n",
    "    plt.pause(0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "overlay_pred_true()\n",
    "overlay_pred_true()\n",
    "overlay_pred_true()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Viewing effect of different final $\\Delta t$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def viz_diff_t():\n",
    "    random_index = np.random.choice(len(rand_x))\n",
    "    colors = ['Reds', 'Blues', 'Greens']\n",
    "    alphas = [0.5, 0.3, 0.3]\n",
    "    for i in range(1, 4):\n",
    "        rand_t[:, -2:-1, :, :, :] = torch.ones_like(rand_t[:, -2:-1, :, :, :])*i/10\n",
    "        rand_y_hat_alt = conv_time_lstm(rand_x.to(device),\n",
    "                                        rand_t.to(device))[0][0][:, -2:-1, :, :, :]\n",
    "        rand_y_hat_alt = rand_y_hat_alt.cpu().data.numpy()\n",
    "        plt.imshow(rand_y_hat_alt[random_index, 0, 0], cmap = colors[i-1], alpha = alphas[i-1])\n",
    "    plt.pause(0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "viz_diff_t()\n",
    "viz_diff_t()\n",
    "viz_diff_t()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mimicking sequence-to-sequence by shuffling in predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "random_index = np.random.choice(len(rand_x) - 1)\n",
    "rand_x = rand_x[random_index:(random_index+1)]\n",
    "\n",
    "for i in range(10):\n",
    "    f, axarr = plt.subplots(1, 2)\n",
    "    f.set_figheight(3)\n",
    "    f.set_figwidth(6)\n",
    "    axarr[0].imshow(rand_x[0, 0, 0], cmap = 'gist_gray')\n",
    "    rand_y_hat = conv_time_lstm(rand_x.to(device),\n",
    "                                torch.ones_like(rand_x.to(device)) / 10)[0][0][:, -2:-1, :, :, :]\n",
    "    axarr[1].imshow(rand_y_hat[0, 0, 0].data.cpu().numpy(), cmap = 'Blues')\n",
    "    rand_x = torch.cat([rand_x, rand_y_hat.data.cpu()], dim = 1)\n",
    "    rand_x = rand_x[:, 1:]\n",
    "    plt.pause(0.01)"
   ]
  },
  {
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}