colab_train.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "train.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3TvbK886gHnZ"
      },
      "source": [
        "# Mount Drive"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xSSIRR1PPnHS"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WUcJtigDTlGK"
      },
      "source": [
        "%cd /content/MyDrive/MyDrive/PredNet_ApparentMotion"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NPNYOGTKgOV1"
      },
      "source": [
        "# Import"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cbdZL2JnPxyL"
      },
      "source": [
        "!pip install torchinfo\n",
        "!pip install hickle"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RWYxojm2P03t"
      },
      "source": [
        "import os\n",
        "import gc\n",
        "import sys\n",
        "import math\n",
        "import numpy as np\n",
        "import hickle as hkl\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.autograd import Variable\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "from torchinfo import summary\n",
        "\n",
        "from prednet import PredNet"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-0JlgsgiP36V"
      },
      "source": [
        "gc.collect()\n",
        "torch.cuda.empty_cache()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ymmwh8PygRll"
      },
      "source": [
        "# Setting Variables and Parameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u44bbbl5P5BJ"
      },
      "source": [
        "num_epochs = 150\n",
        "batch_size = 5\n",
        "A_channels = (3, 48, 96, 192)\n",
        "R_channels = (3, 48, 96, 192)\n",
        "lr = 0.001 # if epoch < 75 0.001 else 0.0001\n",
        "nt = 10 # num of time steps"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XTbgsxC8PtDX"
      },
      "source": [
        "layer_loss_weights = Variable(torch.FloatTensor([[1.], [0.], [0.], [0.]]).cuda())\n",
        "time_loss_weights = 1./(nt - 1) * torch.ones(nt, 1)\n",
        "time_loss_weights[0] = 0\n",
        "time_loss_weights = Variable(time_loss_weights.cuda())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "apM6C5-hP9hQ"
      },
      "source": [
        "train_file = './kitti_data_raw/X_train.hkl'\n",
        "train_sources = './kitti_data_raw/sources_train.hkl'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m0yGbEo2QDog"
      },
      "source": [
        "X = hkl.load(train_file)\n",
        "step = nt * batch_size\n",
        "portion = math.floor(X.shape[0]/step)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v1iWrOqEgVmD"
      },
      "source": [
        "# Set Model and Optimizer"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3s6s8l7nP_Vg"
      },
      "source": [
        "model = PredNet(R_channels, A_channels, output_mode='error')\n",
        "if torch.cuda.is_available():\n",
        "    print('Using GPU.')\n",
        "    model.cuda()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "maEkI24kQBs7"
      },
      "source": [
        "optimizer = torch.optim.SGD(model.parameters(), lr=lr)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jEbrtBJdQCjC"
      },
      "source": [
        "def lr_scheduler(optimizer, epoch):\n",
        "    if epoch < num_epochs //2:\n",
        "        return optimizer\n",
        "    else:\n",
        "        for param_group in optimizer.param_groups:\n",
        "            param_group['lr'] = 0.0001\n",
        "        return optimizer"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "11QmZxZZQGwA"
      },
      "source": [
        "def makeinputs(inputs):\n",
        "    X = []\n",
        "    ind = 0\n",
        "    for b in range(batch_size):\n",
        "        temp_batch = []\n",
        "        for n in range(nt):\n",
        "            temp_batch.append(inputs[ind].transpose(2,0,1).tolist())\n",
        "            ind += 1\n",
        "        X.append(temp_batch)\n",
        "    return Variable(torch.Tensor(X).cuda())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C-Q5qOwsgn3B"
      },
      "source": [
        "# Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c_gVTbh1QIZx"
      },
      "source": [
        "temporal_loss = []\n",
        "\n",
        "for epoch in range(num_epochs):\n",
        "    optimizer = lr_scheduler(optimizer, epoch)\n",
        "    loss = 0\n",
        "    for p in range(portion):\n",
        "        Xtrim = X[p*step:(p+1)*step]\n",
        "        inputs = makeinputs(Xtrim)\n",
        "            \n",
        "        loss = model(inputs) # batch x n_layers x nt\n",
        "        \n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        \n",
        "        if (p+1)%10 == 0:\n",
        "            print('step: {}/{}, loss: {}'.format(p+1,portion,loss))\n",
        "            temporal_loss.append(float(loss.data.cpu().detach()))\n",
        "    print('Epoch: {}/{}, loss: {}'.format(epoch+1, num_epochs, loss))\n",
        "    \n",
        "    with torch.no_grad():\n",
        "        torch.cuda.empty_cache()\n",
        "    \n",
        "    if epoch%10 == 0:\n",
        "        torch.save(model.state_dict(), './models/training_{:04d}.pt'.format(epoch))\n",
        "\n",
        "torch.save(model.state_dict(), './models/training.pt')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}