09_prepare_for_batch_training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from datetime import datetime\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import torch\n",
    "from torch.autograd import Variable\n",
    "import torch.nn as nn\n",
    "from torch.utils import data\n",
    "from models.ConvTimeLSTM2 import ConvTime_LSTM2\n",
    "from helper_fns.processing import scale_and_remove_na"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br><br><br><br>\n",
    "# Required input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_input_scenes = 10\n",
    "train_percent = 0.70\n",
    "out_samp_perc = 0.15 # validation and testing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basic data import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(49, 10, 5, 96, 96)\n",
      "(151, 10, 5, 96, 96)\n",
      "(49, 10, 5, 96, 96)\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-3-866d31a4894f>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     12\u001b[0m     \u001b[0mnumpy_data_location\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"data/\"\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mvol\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m\"/numpy_data_cube.npy\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m     \u001b[0mtable_data_location\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"data/\"\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mvol\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m\"/good_df.csv\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m     \u001b[0mvolcano_scenes\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnumpy_data_location\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     15\u001b[0m     \u001b[0mtabular_metadata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mread_csv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtable_data_location\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/numpy/lib/npyio.py\u001b[0m in \u001b[0;36mload\u001b[0;34m(file, mmap_mode, allow_pickle, fix_imports, encoding)\u001b[0m\n\u001b[1;32m    451\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    452\u001b[0m                 return format.read_array(fid, allow_pickle=allow_pickle,\n\u001b[0;32m--> 453\u001b[0;31m                                          pickle_kwargs=pickle_kwargs)\n\u001b[0m\u001b[1;32m    454\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    455\u001b[0m             \u001b[0;31m# Try a pickle\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/numpy/lib/format.py\u001b[0m in \u001b[0;36mread_array\u001b[0;34m(fp, allow_pickle, pickle_kwargs)\u001b[0m\n\u001b[1;32m    753\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0misfileobj\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfp\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    754\u001b[0m             \u001b[0;31m# We can use the fast fromfile() function.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 755\u001b[0;31m             \u001b[0marray\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnumpy\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfromfile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfp\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcount\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcount\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    756\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    757\u001b[0m             \u001b[0;31m# This is not a real file. We have to read it the\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "volcanoes = os.listdir(\"data\")\n",
    "\n",
    "try:\n",
    "    volcanoes.remove(\".ipynb_checkpoints\")\n",
    "except ValueError as e:\n",
    "    do = 'nothing'\n",
    "    \n",
    "count = 0\n",
    "for vol in volcanoes:\n",
    "    \n",
    "    ### Basic data import ###\n",
    "    numpy_data_location = \"data/\" + vol + \"/numpy_data_cube.npy\"\n",
    "    table_data_location = \"data/\" + vol + \"/good_df.csv\"\n",
    "    volcano_scenes = np.load(numpy_data_location)\n",
    "    tabular_metadata = pd.read_csv(table_data_location)\n",
    "    \n",
    "    ### Separate model inputs and outputs\n",
    "    # Determine number in each partition\n",
    "    train_n = int(np.floor((len(volcano_scenes) - num_input_scenes)*train_percent))\n",
    "    out_n = int(np.floor((len(volcano_scenes) - num_input_scenes)*out_samp_perc))\n",
    "\n",
    "    # For every data partition\n",
    "    # Array for the prior scenes\n",
    "    x_scenes_train = np.zeros([train_n,\n",
    "                               num_input_scenes,\n",
    "                               volcano_scenes.shape[1],\n",
    "                               volcano_scenes.shape[2],\n",
    "                               volcano_scenes.shape[3]])\n",
    "    x_scenes_valid = np.zeros([out_n,\n",
    "                               num_input_scenes,\n",
    "                               volcano_scenes.shape[1],\n",
    "                               volcano_scenes.shape[2],\n",
    "                               volcano_scenes.shape[3]])\n",
    "    x_scenes_test = np.zeros([out_n,\n",
    "                              num_input_scenes,\n",
    "                              volcano_scenes.shape[1],\n",
    "                              volcano_scenes.shape[2],\n",
    "                              volcano_scenes.shape[3]])\n",
    "    # Array for the time differences between scenes\n",
    "    time_differences_train = np.ones(x_scenes_train.shape)\n",
    "    time_differences_valid = np.ones(x_scenes_valid.shape)\n",
    "    time_differences_test = np.ones(x_scenes_test.shape)\n",
    "    # Array for the target scenes\n",
    "    y_scenes_train = np.zeros([train_n,\n",
    "                               1,\n",
    "                               volcano_scenes.shape[1],\n",
    "                               volcano_scenes.shape[2],\n",
    "                               volcano_scenes.shape[3]])\n",
    "    y_scenes_valid = np.zeros([out_n,\n",
    "                               1,\n",
    "                               volcano_scenes.shape[1],\n",
    "                               volcano_scenes.shape[2],\n",
    "                               volcano_scenes.shape[3]])\n",
    "    y_scenes_test = np.zeros([out_n,\n",
    "                              1,\n",
    "                              volcano_scenes.shape[1],\n",
    "                              volcano_scenes.shape[2],\n",
    "                              volcano_scenes.shape[3]])\n",
    "    # Array for the prior max temperature above the background\n",
    "    x_temperatures_train = np.zeros([train_n,\n",
    "                                     num_input_scenes])\n",
    "    x_temperatures_valid = np.zeros([out_n,\n",
    "                                     num_input_scenes])\n",
    "    x_temperatures_test = np.zeros([out_n,\n",
    "                                    num_input_scenes])\n",
    "    # Array for the target max temperature above the background\n",
    "    y_temperatures_train = np.zeros([train_n])\n",
    "    y_temperatures_valid = np.zeros([out_n])\n",
    "    y_temperatures_test = np.zeros([out_n])\n",
    "    # Formatting the string dates as datetime objects\n",
    "    formatted_dates = [datetime.strptime(date, '%Y-%m-%d') for date in tabular_metadata['dates']]\n",
    "    # For all observations - acknowledging that the first (n-1) wont have n prior observations\n",
    "    for i in range(num_input_scenes, x_scenes_train.shape[0] + x_scenes_valid.shape[0] + x_scenes_test.shape[0] + 10):\n",
    "        if i < (train_n + num_input_scenes):\n",
    "            # Store the image data\n",
    "            x_scenes_train[i - num_input_scenes, :, :, :, :] = volcano_scenes[(i - num_input_scenes):i, :, :, :]\n",
    "            y_scenes_train[i - num_input_scenes, 0, :, :, :] = volcano_scenes[i, :, :, :]\n",
    "            # Store the max temperature scalars\n",
    "            x_temperatures_train[i - num_input_scenes, :] = tabular_metadata['T_above_back'].values[(i - num_input_scenes):i]\n",
    "            y_temperatures_train[i - num_input_scenes] = tabular_metadata['T_above_back'].values[i]\n",
    "            # Compute the time differences and store\n",
    "            dates_i_plus_1 = formatted_dates[(i - num_input_scenes + 1):(i + 1)]\n",
    "            dates_i = formatted_dates[(i - num_input_scenes):i]\n",
    "            for j in range(len(dates_i_plus_1)):\n",
    "                time_differences_train[i - num_input_scenes, j] = (dates_i_plus_1[j] - dates_i[j]).days\n",
    "        elif i < (train_n + out_n + num_input_scenes):\n",
    "            # Store the image data\n",
    "            x_scenes_valid[i - train_n - num_input_scenes, :, :, :, :] = volcano_scenes[(i - num_input_scenes):i, :, :, :]\n",
    "            y_scenes_valid[i - train_n - num_input_scenes, 0, :, :, :] = volcano_scenes[i, :, :, :]\n",
    "            # Store the max temperature scalars\n",
    "            x_temperatures_valid[i - train_n - num_input_scenes, :] = tabular_metadata['T_above_back'].values[(i - num_input_scenes):i]\n",
    "            y_temperatures_valid[i - train_n - num_input_scenes] = tabular_metadata['T_above_back'].values[i]\n",
    "            # Compute the time differences and store\n",
    "            dates_i_plus_1 = formatted_dates[(i - num_input_scenes + 1):(i + 1)]\n",
    "            dates_i = formatted_dates[(i - num_input_scenes):i]\n",
    "            for j in range(len(dates_i_plus_1)):\n",
    "                time_differences_valid[i - train_n - num_input_scenes, j] = (dates_i_plus_1[j] - dates_i[j]).days\n",
    "        else:\n",
    "            # Store the image data\n",
    "            x_scenes_test[i - train_n - out_n - num_input_scenes, :, :, :, :] = volcano_scenes[(i - num_input_scenes):i, :, :, :]\n",
    "            y_scenes_test[i - train_n - out_n - num_input_scenes, 0, :, :, :] = volcano_scenes[i, :, :, :]\n",
    "            # Store the max temperature scalars\n",
    "            x_temperatures_test[i - train_n - out_n - num_input_scenes, :] = tabular_metadata['T_above_back'].values[(i - num_input_scenes):i]\n",
    "            y_temperatures_test[i - train_n - out_n - num_input_scenes] = tabular_metadata['T_above_back'].values[i]\n",
    "            # Compute the time differences and store\n",
    "            dates_i_plus_1 = formatted_dates[(i - num_input_scenes + 1):(i + 1)]\n",
    "            dates_i = formatted_dates[(i - num_input_scenes):i]\n",
    "            for j in range(len(dates_i_plus_1)):\n",
    "                time_differences_test[i - train_n - out_n - num_input_scenes, j] = (dates_i_plus_1[j] - dates_i[j]).days\n",
    "    \n",
    "    if count == 0:\n",
    "        x_train = x_scenes_train\n",
    "        t_train = time_differences_train\n",
    "        y_train = y_scenes_train\n",
    "        x_valid = x_scenes_valid\n",
    "        t_valid = time_differences_valid\n",
    "        y_valid = y_scenes_valid\n",
    "        x_test = x_scenes_test\n",
    "        t_test = time_differences_test\n",
    "        y_test = y_scenes_test\n",
    "    else:\n",
    "        x_train = np.append(x_train, x_scenes_train, axis = 0)\n",
    "        t_train = np.append(t_train, time_differences_train, axis = 0)\n",
    "        y_train = np.append(y_train, y_scenes_train, axis = 0)\n",
    "        x_valid = np.append(x_valid, x_scenes_valid, axis = 0)\n",
    "        t_valid = np.append(t_valid, time_differences_valid, axis = 0)\n",
    "        y_valid = np.append(y_valid, y_scenes_valid, axis = 0)\n",
    "        x_test = np.append(x_test, x_scenes_test, axis = 0)\n",
    "        t_test = np.append(t_test, time_differences_test, axis = 0)\n",
    "        y_test = np.append(y_test, y_scenes_train, axis = 0)\n",
    "\n",
    "    count += 1\n",
    "    \n",
    "    \n",
    "    print(x_scenes_train.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(200, 10, 5, 96, 96)"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.append(np.zeros([49, 10, 5, 96, 96]),\n",
    "           np.zeros([151, 10, 5, 96, 96]),\n",
    "          axis = 0).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Help on function append in module numpy:\n",
      "\n",
      "append(arr, values, axis=None)\n",
      "    Append values to the end of an array.\n",
      "    \n",
      "    Parameters\n",
      "    ----------\n",
      "    arr : array_like\n",
      "        Values are appended to a copy of this array.\n",
      "    values : array_like\n",
      "        These values are appended to a copy of `arr`.  It must be of the\n",
      "        correct shape (the same shape as `arr`, excluding `axis`).  If\n",
      "        `axis` is not specified, `values` can be any shape and will be\n",
      "        flattened before use.\n",
      "    axis : int, optional\n",
      "        The axis along which `values` are appended.  If `axis` is not\n",
      "        given, both `arr` and `values` are flattened before use.\n",
      "    \n",
      "    Returns\n",
      "    -------\n",
      "    append : ndarray\n",
      "        A copy of `arr` with `values` appended to `axis`.  Note that\n",
      "        `append` does not occur in-place: a new array is allocated and\n",
      "        filled.  If `axis` is None, `out` is a flattened array.\n",
      "    \n",
      "    See Also\n",
      "    --------\n",
      "    insert : Insert elements into an array.\n",
      "    delete : Delete elements from an array.\n",
      "    \n",
      "    Examples\n",
      "    --------\n",
      "    >>> np.append([1, 2, 3], [[4, 5, 6], [7, 8, 9]])\n",
      "    array([1, 2, 3, ..., 7, 8, 9])\n",
      "    \n",
      "    When `axis` is specified, `values` must have the correct shape.\n",
      "    \n",
      "    >>> np.append([[1, 2, 3], [4, 5, 6]], [[7, 8, 9]], axis=0)\n",
      "    array([[1, 2, 3],\n",
      "           [4, 5, 6],\n",
      "           [7, 8, 9]])\n",
      "    >>> np.append([[1, 2, 3], [4, 5, 6]], [7, 8, 9], axis=0)\n",
      "    Traceback (most recent call last):\n",
      "        ...\n",
      "    ValueError: all the input arrays must have same number of dimensions\n",
      "\n"
     ]
    }
   ],
   "source": [
    "help(np.append)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scale temperatures between 0 and 1. If temperature is missing, assigned a scaled value of 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_scenes_train = scale_and_remove_na(x_scenes_train)\n",
    "x_scenes_train = scale_and_remove_na(x_scenes_train)\n",
    "x_scenes_train = scale_and_remove_na(x_scenes_train)\n",
    "\n",
    "time_differences_train = scale_and_remove_na(time_differences_train)\n",
    "time_differences_train = scale_and_remove_na(time_differences_train)\n",
    "time_differences_train = scale_and_remove_na(time_differences_train)\n",
    "\n",
    "y_scenes_train = scale_and_remove_na(y_scenes_train)\n",
    "y_scenes_train = scale_and_remove_na(y_scenes_train)\n",
    "y_scenes_train = scale_and_remove_na(y_scenes_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Passing to pytorch and formatting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_scenes_train = torch.from_numpy(x_scenes_train).type(torch.FloatTensor)\n",
    "x_scenes_test = torch.from_numpy(x_scenes_test).type(torch.FloatTensor)\n",
    "x_scenes_valid = torch.from_numpy(x_scenes_valid).type(torch.FloatTensor)\n",
    "\n",
    "time_differences_train = torch.from_numpy(time_differences_train).type(torch.FloatTensor)\n",
    "time_differences_test = torch.from_numpy(time_differences_test).type(torch.FloatTensor)\n",
    "time_differences_valid = torch.from_numpy(time_differences_valid).type(torch.FloatTensor)\n",
    "\n",
    "y_scenes_train = torch.from_numpy(y_scenes_train).type(torch.FloatTensor)\n",
    "y_scenes_test = torch.from_numpy(y_scenes_test).type(torch.FloatTensor)\n",
    "y_scenes_valid = torch.from_numpy(y_scenes_valid).type(torch.FloatTensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining model parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Picking one of the like-sequence tensors within the list to set parameters\n",
    "channels = x_scenes_train.shape[2]\n",
    "height = x_scenes_train.shape[3]\n",
    "width = x_scenes_train.shape[4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "conv_time_lstm = ConvTime_LSTM2(input_size = (height,\n",
    "                                              width),\n",
    "                                input_dim = channels,\n",
    "                                hidden_dim = [128, 64, 64, 1],\n",
    "                                kernel_size = (5, 5),\n",
    "                                num_layers = 4,\n",
    "                                batch_first = True,\n",
    "                                bias = True,\n",
    "                                return_all_layers = False,\n",
    "                                GPU = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Passing to GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ConvTime_LSTM2(\n",
       "  (cell_list): ModuleList(\n",
       "    (0): ConvTime_LSTM2Cell(\n",
       "      (i_conv): Conv2d(133, 128, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (T1_conv_x): Conv2d(5, 128, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (T1_conv_t): Conv2d(5, 128, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (T2_conv_x): Conv2d(5, 128, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (T2_conv_t): Conv2d(5, 128, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (c_conv): Conv2d(133, 128, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (o_conv): Conv2d(138, 128, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "    )\n",
       "    (1): ConvTime_LSTM2Cell(\n",
       "      (i_conv): Conv2d(192, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (T1_conv_x): Conv2d(128, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (T1_conv_t): Conv2d(128, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (T2_conv_x): Conv2d(128, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (T2_conv_t): Conv2d(128, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (c_conv): Conv2d(192, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (o_conv): Conv2d(320, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "    )\n",
       "    (2): ConvTime_LSTM2Cell(\n",
       "      (i_conv): Conv2d(128, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (T1_conv_x): Conv2d(64, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (T1_conv_t): Conv2d(64, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (T2_conv_x): Conv2d(64, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (T2_conv_t): Conv2d(64, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (c_conv): Conv2d(128, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (o_conv): Conv2d(192, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "    )\n",
       "    (3): ConvTime_LSTM2Cell(\n",
       "      (i_conv): Conv2d(65, 1, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (T1_conv_x): Conv2d(64, 1, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (T1_conv_t): Conv2d(64, 1, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (T2_conv_x): Conv2d(64, 1, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (T2_conv_t): Conv2d(64, 1, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (c_conv): Conv2d(65, 1, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "      (o_conv): Conv2d(129, 1, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv_time_lstm.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting optimization methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss = torch.nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(conv_time_lstm.parameters())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining data set and data loaders for parallelization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "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_scenes_train[IDs, :, :, :, :]\n",
    "        curr_t = time_differences_train[IDs, :, :, :, :]\n",
    "        curr_y = y_scenes_train[IDs, :, :, :, :]\n",
    "\n",
    "        #return X, y\n",
    "        return(curr_x, curr_t, curr_y)\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_scenes_valid[IDs, :, :, :, :]\n",
    "        curr_t = time_differences_valid[IDs, :, :, :, :]\n",
    "        curr_y = y_scenes_valid[IDs, :, :, :, :]\n",
    "\n",
    "        #return X, y\n",
    "        return(curr_x, curr_t, curr_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_set = train_Dataset(data_indices=range(y_scenes_train.shape[0]))\n",
    "validation_set = validation_Dataset(data_indices=range(y_scenes_valid.shape[0]))\n",
    "\n",
    "batch_size = 16\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": "markdown",
   "metadata": {},
   "source": [
    "## Retrieving available computing devices and using parallel GPUs if available"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "conv_time_lstm = torch.nn.DataParallel(conv_time_lstm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 138, 96, 96])\n",
      "torch.Size([16, 138, 96, 96])\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "CUDA out of memory. Tried to allocate 72.00 MiB (GPU 0; 3.95 GiB total capacity; 3.07 GiB already allocated; 29.88 MiB free; 3.25 GiB reserved in total by PyTorch)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-21-f2c78ebfb791>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     14\u001b[0m         \u001b[0;31m# run model and get the prediction\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m         batch_y_hat = conv_time_lstm(batch_x,\n\u001b[0;32m---> 16\u001b[0;31m                                      batch_t)\n\u001b[0m\u001b[1;32m     17\u001b[0m         \u001b[0mbatch_y_hat\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbatch_y_hat\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m-\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    548\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    549\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 550\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    551\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    552\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, *inputs, **kwargs)\u001b[0m\n\u001b[1;32m    151\u001b[0m         \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mscatter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice_ids\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    152\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice_ids\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 153\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    154\u001b[0m         \u001b[0mreplicas\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreplicate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodule\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice_ids\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    155\u001b[0m         \u001b[0moutputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparallel_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mreplicas\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    548\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    549\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 550\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    551\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    552\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/VolcanoesIR/models/ConvTimeLSTM2.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input_tensor, time_tensor, hidden_state)\u001b[0m\n\u001b[1;32m    245\u001b[0m                 h, c = self.cell_list[layer_idx](input_tensor = cur_layer_input[:, t, :, :, :],\n\u001b[1;32m    246\u001b[0m                                                  \u001b[0mtime_tensor\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcur_time_input\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 247\u001b[0;31m                                                  cur_state=[h, c])\n\u001b[0m\u001b[1;32m    248\u001b[0m                 \u001b[0moutput_inner\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mh\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    249\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    548\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    549\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 550\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    551\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    552\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/VolcanoesIR/models/ConvTimeLSTM2.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input_tensor, time_tensor, cur_state)\u001b[0m\n\u001b[1;32m    155\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    156\u001b[0m         \u001b[0;31m## The hidden vector ##\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 157\u001b[0;31m         \u001b[0mh_m\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mo_m\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtanh\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mc_m_tilde\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    158\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    159\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: CUDA out of memory. Tried to allocate 72.00 MiB (GPU 0; 3.95 GiB total capacity; 3.07 GiB already allocated; 29.88 MiB free; 3.25 GiB reserved in total by PyTorch)"
     ]
    }
   ],
   "source": [
    "loss_list = []\n",
    "epochs = int(np.ceil((7*10**5) / x_scenes_train.shape[0]))\n",
    "for i in range(epochs):\n",
    "    for data in train_loader:\n",
    "        \n",
    "        # data loader\n",
    "        batch_x, batch_t, batch_y = 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')\n",
    "        \n",
    "    print('Epoch: ', i, '\\n\\tBatch loss: ', batch_loss.item(), '\\n')"
   ]
  },
  {
   "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}