Main_Train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import shutil\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable\n",
    "\n",
    "from hybrid_net import DataTransformer, DataLoaderMaker, ModelArchitecture, MetricMeters, HyperParameterSchedulers, Losses\n",
    "\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Important Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = \"./data-local/images/cifar10/by-image/\"\n",
    "labels = \"./data-local/labels/cifar10/1000_balanced_labels/00.txt\"\n",
    "train_subdir = \"train\"\n",
    "eval_subdir = \"val\"\n",
    "checkpoint_path = \"./checkpoints/\"\n",
    "device = \"cuda\" if torch.cuda.is_available() else 'cpu'\n",
    "workers = 4\n",
    "NO_LABEL = -1\n",
    "global_step = 0\n",
    "ema_decay = 0.999\n",
    "print_freq = 2\n",
    "best_prec1 = 0\n",
    "start_epoch = 0\n",
    "arch = \"hybridNet1\"\n",
    "\n",
    "\n",
    "# specified in the paper\n",
    "initial_lr = 0.003\n",
    "initial_beta1 = 0.9\n",
    "\n",
    "# batch_size is 100 in paper\n",
    "batch_size = 10\n",
    "\n",
    "# batch_size is 20 in paper\n",
    "labeled_batch_size = 1\n",
    "\n",
    "# total epochs not specified for 1000 label CIFAR training in paper\n",
    "total_epochs = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function to update the secondary model (used for stability loss)\n",
    "\n",
    "def update_ema_variables(model, ema_model, alpha, global_step):\n",
    "    # Use the true average until the exponential average is more correct\n",
    "    alpha = min(1 - 1 / (global_step + 1), alpha)\n",
    "    for ema_param, param in zip(ema_model.parameters(), model.parameters()):\n",
    "        ema_param.data.mul_(alpha).add_(1 - alpha, param.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function to save checkpoints of the model after every epoch\n",
    "\n",
    "def save_checkpoint(state, is_best, dirpath, epoch):\n",
    "    filename = 'checkpoint-{}.ckpt'.format(epoch)\n",
    "    checkpoint_path = os.path.join(dirpath, filename)\n",
    "    best_path = os.path.join(dirpath, 'best.ckpt')\n",
    "    torch.save(state, checkpoint_path)\n",
    "    print('--- checkpoint saved to %s ---'.format(checkpoint_path))\n",
    "    if is_best:\n",
    "        shutil.copyfile(checkpoint_path, best_path)\n",
    "        print('--- checkpoint copied to %s ---'.format(best_path))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DataLoader Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_transformation, eval_transformation = DataTransformer.transformer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader, eval_loader = DataLoaderMaker.create_data_loaders(train_transformation, \n",
    "                                                eval_transformation, \n",
    "                                                images, \n",
    "                                                train_subdir, \n",
    "                                                eval_subdir,\n",
    "                                                labels,\n",
    "                                                batch_size,\n",
    "                                                labeled_batch_size,\n",
    "                                                workers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 'model' is the actual model to be trained\n",
    "# 'ema_model' is a duplicate model for stability loss as mentioned in paper\n",
    "\n",
    "model = ModelArchitecture.create_model()\n",
    "ema_model = ModelArchitecture.create_model(ema=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train and Validation Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(train_loader, model, ema_model, optimizer, epoch, ema_decay, total_epochs, print_freq = 2):\n",
    "    global global_step\n",
    "    \n",
    "    \n",
    "    # initialize loss functions\n",
    "    class_criterion = nn.CrossEntropyLoss(size_average=False, ignore_index=NO_LABEL).to(device)\n",
    "    reconstruction_criterion = Losses.symmetric_mse_loss   \n",
    "    stability_criterion = Losses.softmax_mse_loss\n",
    "    \n",
    "    # initialize evaluation metric meter\n",
    "    meters = MetricMeters.AverageMeterSet()\n",
    "\n",
    "    # switch to train mode\n",
    "    model.train()\n",
    "    ema_model.train()\n",
    "\n",
    "    end = time.time()\n",
    "    \n",
    "    for i, ((input, ema_input), target) in enumerate(train_loader):\n",
    "        # measure data loading time\n",
    "        meters.update('data_time', time.time() - end)\n",
    "        \n",
    "        # update hyperparameters according to details given in the paper\n",
    "        HyperParameterSchedulers.adjust_learning_rate(optimizer, initial_lr, epoch, total_epochs)\n",
    "        HyperParameterSchedulers.adjust_beta_1(optimizer, initial_beta1, epoch, total_epochs)\n",
    "        if epoch == 0:\n",
    "            lambda_c = HyperParameterSchedulers.exponential_increase(i, 800)\n",
    "        else: \n",
    "            lambda_c = 1.0\n",
    "        lambda_r = 100 * HyperParameterSchedulers.adjust_lambda_r(epoch, 0.25 * total_epochs, 0.8 * total_epochs, total_epochs)\n",
    "        lambda_s = HyperParameterSchedulers.exponential_decrease(epoch, 0.95 * total_epochs, total_epochs)\n",
    "        meters.update('lr', optimizer.param_groups[0]['lr'])\n",
    "        meters.update('beta1', optimizer.param_groups[0]['betas'][0])\n",
    "        meters.update('lambda_c', lambda_c)\n",
    "        meters.update('lambda_r', lambda_r)\n",
    "        meters.update('lambda_s', lambda_s)\n",
    "        \n",
    "        # prepare input and target \n",
    "        input_var = torch.autograd.Variable(input)\n",
    "        ema_input_var = torch.autograd.Variable(ema_input, volatile=True)\n",
    "        target_var = torch.autograd.Variable(target.cuda(async=True))\n",
    "        \n",
    "        input_var = input_var.to(device)\n",
    "        ema_input_var = ema_input_var.to(device)\n",
    "\n",
    "        minibatch_size = len(target_var)\n",
    "        labeled_minibatch_size = target_var.data.ne(NO_LABEL).sum().float()\n",
    "        \n",
    "        assert labeled_minibatch_size > 0\n",
    "        meters.update('labeled_minibatch_size', labeled_minibatch_size)\n",
    "        \n",
    "        \n",
    "        # forward pass through two models\n",
    "        ema_y, ema_x_c, ema_x_u = ema_model(ema_input_var)\n",
    "        model_y, model_x_c, model_x_u = model(input_var)\n",
    "\n",
    "        # get logits of outputs from the model\n",
    "        ema_logit = Variable(ema_y.detach().data, requires_grad=False)\n",
    "        cons_logit = model_y\n",
    "        class_logit = model_y\n",
    "        \n",
    "        # classification loss\n",
    "        class_loss = class_criterion(class_logit, target_var) / labeled_minibatch_size\n",
    "        meters.update('class_loss', class_loss.data[0])\n",
    "        \n",
    "        ema_class_loss = class_criterion(ema_logit, target_var) / labeled_minibatch_size\n",
    "        meters.update('ema_class_loss', ema_class_loss.data[0])\n",
    "\n",
    "        # reconstruction loss\n",
    "        reconstruction_loss = reconstruction_criterion(model_x_c + model_x_u, input_var)/minibatch_size\n",
    "        meters.update('reconstruction_loss', reconstruction_loss)\n",
    "        \n",
    "        # TODO add reconstruction loss of in-between layers\n",
    "        #\n",
    "        #\n",
    "        #\n",
    "        reconstruction_loss_bl = 0\n",
    "        \n",
    "        # balanced reconstruction loss\n",
    "        if torch.sum((model_x_c - input_var)**2) <= torch.sum((model_x_u - input_var)**2):\n",
    "            reconstruction_loss = reconstruction_criterion(model_x_u + model_x_c.detach(), input_var)/minibatch_size\n",
    "        else:\n",
    "            reconstruction_loss = reconstruction_criterion(model_x_u.detach() + model_x_c, input_var)/minibatch_size\n",
    "      \n",
    "        \n",
    "        # stability loss\n",
    "        stability_loss = stability_criterion(cons_logit, ema_logit) / minibatch_size\n",
    "        meters.update('stability_loss', stability_loss.data[0])\n",
    "\n",
    "\n",
    "        # final loss as given in paper\n",
    "        loss = lambda_c * class_loss + lambda_r * reconstruction_loss + lambda_r * reconstruction_loss_bl + lambda_s * stability_loss \n",
    "        \n",
    "        # ensure the loss is not exploding\n",
    "        assert not (np.isnan(loss.data[0]) or loss.data[0] > 1e5), 'Loss explosion: {}'.format(loss.data[0])\n",
    "        \n",
    "        meters.update('loss', loss.data[0])\n",
    "\n",
    "        # evaluate performance of model after training\n",
    "        prec1, prec5 = MetricMeters.accuracy(class_logit.data, target_var.data, topk=(1, 5))\n",
    "        meters.update('top1', prec1[0], labeled_minibatch_size)\n",
    "        meters.update('error1', 100. - prec1[0], labeled_minibatch_size)\n",
    "        meters.update('top5', prec5[0], labeled_minibatch_size)\n",
    "        meters.update('error5', 100. - prec5[0], labeled_minibatch_size)\n",
    "        \n",
    "        # evaluate performance of ema_model after training\n",
    "        ema_prec1, ema_prec5 = MetricMeters.accuracy(ema_logit.data, target_var.data, topk=(1, 5))\n",
    "        meters.update('ema_top1', ema_prec1[0], labeled_minibatch_size)\n",
    "        meters.update('ema_error1', 100. - ema_prec1[0], labeled_minibatch_size)\n",
    "        meters.update('ema_top5', ema_prec5[0], labeled_minibatch_size)\n",
    "        meters.update('ema_error5', 100. - ema_prec5[0], labeled_minibatch_size)\n",
    "\n",
    "        # compute gradient and do Adam step\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        global_step += 1\n",
    "        \n",
    "        # update parameters of 'ema_model' as it is not trained\n",
    "        update_ema_variables(model, ema_model, ema_decay, global_step)\n",
    "\n",
    "        # measure elapsed time\n",
    "        meters.update('batch_time', time.time() - end)\n",
    "        end = time.time()\n",
    "\n",
    "        # print training evaluation metrics\n",
    "        if i % print_freq == 0:\n",
    "            print(\n",
    "                'Epoch: [{0}][{1}/{2}]\\t'\n",
    "                'Class {meters[class_loss]:.4f}\\t'\n",
    "                'Prec@1 {meters[top1]:.3f}\\t'\n",
    "                'Prec@5 {meters[top5]:.3f}'.format(\n",
    "                    epoch, i, len(train_loader), meters=meters))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def validate(eval_loader, model, global_step, epoch, print_freq = 2):\n",
    "    \n",
    "    # initialize loss functions\n",
    "    class_criterion = nn.CrossEntropyLoss(size_average=False, ignore_index=NO_LABEL).to(device)\n",
    "    \n",
    "    # initialize evaluation metric meter\n",
    "    meters = MetricMeters.AverageMeterSet()\n",
    "\n",
    "    # switch to evaluate mode\n",
    "    model.eval()\n",
    "\n",
    "    end = time.time()\n",
    "    for i, (input, target) in enumerate(eval_loader):\n",
    "        # measure data loading time\n",
    "        meters.update('data_time', time.time() - end)\n",
    "        \n",
    "        \n",
    "        # prepare input and target\n",
    "        input_var = torch.autograd.Variable(input, volatile=True)\n",
    "        target_var = torch.autograd.Variable(target.cuda(async=True), volatile=True)\n",
    "\n",
    "        input_var = input_var.to(device)\n",
    "\n",
    "        minibatch_size = len(target_var)\n",
    "        labeled_minibatch_size = target_var.data.ne(NO_LABEL).sum().float()\n",
    "        assert labeled_minibatch_size > 0\n",
    "        meters.update('labeled_minibatch_size', labeled_minibatch_size)\n",
    "\n",
    "        # forward pass through the model\n",
    "        model_y, model_x_c, model_x_u = model(input_var)\n",
    "        softmax_model_y = F.softmax(model_y, dim=1)\n",
    "        class_loss = class_criterion(model_y, target_var) / minibatch_size\n",
    "\n",
    "        # measure accuracy and record loss\n",
    "        prec1, prec5 = MetricMeters.accuracy(model_y.data, target_var.data, topk=(1, 5))\n",
    "        meters.update('class_loss', class_loss.data[0], labeled_minibatch_size)\n",
    "        meters.update('top1', prec1[0], labeled_minibatch_size)\n",
    "        meters.update('error1', 100.0 - prec1[0], labeled_minibatch_size)\n",
    "        meters.update('top5', prec5[0], labeled_minibatch_size)\n",
    "        meters.update('error5', 100.0 - prec5[0], labeled_minibatch_size)\n",
    "\n",
    "        # measure elapsed time\n",
    "        meters.update('batch_time', time.time() - end)\n",
    "        end = time.time()\n",
    "\n",
    "        # print evaluation metric for particular batches\n",
    "        if i % print_freq == 0:\n",
    "            print(\n",
    "                'Test: [{0}/{1}]\\t'\n",
    "                'Class {meters[class_loss]:.4f}\\t'\n",
    "                'Prec@1 {meters[top1]:.3f}\\t'\n",
    "                'Prec@5 {meters[top5]:.3f}'.format(\n",
    "                    i, len(eval_loader), meters=meters))\n",
    "\n",
    "    # print overall evaluation metric for the model\n",
    "    print(' * Prec@1 {top1.avg:.3f}\\tPrec@5 {top5.avg:.3f}'\n",
    "          .format(top1=meters['top1'], top5=meters['top5']))\n",
    "\n",
    "    return meters['top1'].avg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Running the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use Adam optimizer as mentioned in the paper\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr = initial_lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.5/dist-packages/torch/nn/functional.py:52: UserWarning: size_average and reduce args will be deprecated, please use reduction='sum' instead.\n",
      "  warnings.warn(warning.format(ret))\n",
      "/usr/local/lib/python3.5/dist-packages/ipykernel_launcher.py:40: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.\n",
      "/usr/local/lib/python3.5/dist-packages/torch/nn/modules/upsampling.py:122: UserWarning: nn.Upsampling is deprecated. Use nn.functional.interpolate instead.\n",
      "  warnings.warn(\"nn.Upsampling is deprecated. Use nn.functional.interpolate instead.\")\n",
      "/usr/local/lib/python3.5/dist-packages/ipykernel_launcher.py:64: UserWarning: invalid index of a 0-dim tensor. This will be an error in PyTorch 0.5. Use tensor.item() to convert a 0-dim tensor to a Python number\n",
      "/usr/local/lib/python3.5/dist-packages/ipykernel_launcher.py:67: UserWarning: invalid index of a 0-dim tensor. This will be an error in PyTorch 0.5. Use tensor.item() to convert a 0-dim tensor to a Python number\n",
      "/usr/local/lib/python3.5/dist-packages/ipykernel_launcher.py:88: UserWarning: invalid index of a 0-dim tensor. This will be an error in PyTorch 0.5. Use tensor.item() to convert a 0-dim tensor to a Python number\n",
      "/usr/local/lib/python3.5/dist-packages/ipykernel_launcher.py:95: UserWarning: invalid index of a 0-dim tensor. This will be an error in PyTorch 0.5. Use tensor.item() to convert a 0-dim tensor to a Python number\n",
      "/usr/local/lib/python3.5/dist-packages/ipykernel_launcher.py:97: UserWarning: invalid index of a 0-dim tensor. This will be an error in PyTorch 0.5. Use tensor.item() to convert a 0-dim tensor to a Python number\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [0][0/4888]\tClass 2.0970 (2.0970)\tPrec@1 100.000 (100.000)\tPrec@5 100.000 (100.000)\n",
      "Epoch: [0][2/4888]\tClass 2.6016 (2.3435)\tPrec@1 0.000 (33.333)\tPrec@5 0.000 (66.667)\n",
      "Epoch: [0][4/4888]\tClass 2.4821 (2.3053)\tPrec@1 0.000 (20.000)\tPrec@5 0.000 (60.000)\n",
      "Epoch: [0][6/4888]\tClass 2.5454 (2.3366)\tPrec@1 0.000 (14.286)\tPrec@5 0.000 (57.143)\n",
      "Epoch: [0][8/4888]\tClass 1.8276 (2.3335)\tPrec@1 100.000 (22.222)\tPrec@5 100.000 (55.556)\n",
      "Epoch: [0][10/4888]\tClass 2.8058 (2.3715)\tPrec@1 0.000 (18.182)\tPrec@5 0.000 (54.545)\n",
      "Epoch: [0][12/4888]\tClass 1.6589 (2.3350)\tPrec@1 100.000 (23.077)\tPrec@5 100.000 (53.846)\n",
      "Epoch: [0][14/4888]\tClass 2.7905 (2.3613)\tPrec@1 0.000 (20.000)\tPrec@5 0.000 (53.333)\n",
      "Epoch: [0][16/4888]\tClass 2.1662 (2.3634)\tPrec@1 0.000 (17.647)\tPrec@5 100.000 (52.941)\n",
      "Epoch: [0][18/4888]\tClass 2.5723 (2.3804)\tPrec@1 0.000 (15.789)\tPrec@5 0.000 (47.368)\n",
      "Epoch: [0][20/4888]\tClass 1.7843 (2.3434)\tPrec@1 100.000 (19.048)\tPrec@5 100.000 (52.381)\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "CUDA error: out of memory",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-11-a77dc007b90f>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      2\u001b[0m     \u001b[0;31m# training\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m     \u001b[0mstart_time\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtime\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m     \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mema_model\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mepoch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mema_decay\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtotal_epochs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprint_freq\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      5\u001b[0m     \u001b[0mtime_elapsed\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtime\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mstart_time\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'epoch training complete in {: .0f}m {:0f}s'\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtime_elapsed\u001b[0m \u001b[0;34m//\u001b[0m \u001b[0;36m60\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtime_elapsed\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0;36m60\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-8-432b3723b78f>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(train_loader, model, ema_model, optimizer, epoch, ema_decay, total_epochs, print_freq)\u001b[0m\n\u001b[1;32m     53\u001b[0m         \u001b[0;31m# forward pass through two models\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     54\u001b[0m         \u001b[0mema_y\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mema_x_c\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mema_x_u\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mema_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mema_input_var\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 55\u001b[0;31m         \u001b[0mmodel_y\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel_x_c\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel_x_u\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_var\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     56\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     57\u001b[0m         \u001b[0;31m# get logits of outputs from the model\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-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    475\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[0m\n\u001b[1;32m    476\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 477\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[0m\n\u001b[0m\u001b[1;32m    478\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[0m\n\u001b[1;32m    479\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[0m\n",
      "\u001b[0;32m~/Desktop/HybridNet/hybrid_net/ModelArchitecture.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input_image)\u001b[0m\n\u001b[1;32m    187\u001b[0m         \u001b[0munpool2u\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munpool2u\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtconv4u\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mindices1u\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    188\u001b[0m         \u001b[0mtconv3u\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mleaky_relu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbntconv3u\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtconv3u\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0munpool2u\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnegative_slope\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0.1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 189\u001b[0;31m         \u001b[0mtconv2u\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mleaky_relu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbntconv2u\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtconv2u\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtconv3u\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnegative_slope\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0.1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    190\u001b[0m         \u001b[0mx_u\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mleaky_relu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbntconv1u\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtconv1u\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtconv2u\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnegative_slope\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0.1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    191\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/torch/nn/functional.py\u001b[0m in \u001b[0;36mleaky_relu\u001b[0;34m(input, negative_slope, inplace)\u001b[0m\n\u001b[1;32m    753\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0minplace\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    754\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_C\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_nn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mleaky_relu_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnegative_slope\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 755\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_C\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_nn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mleaky_relu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnegative_slope\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;34m\u001b[0m\u001b[0m\n\u001b[1;32m    757\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: CUDA error: out of memory"
     ]
    }
   ],
   "source": [
    "for epoch in range(start_epoch, total_epochs):\n",
    "    # training\n",
    "    start_time = time.time()\n",
    "    train(train_loader, model, ema_model, optimizer, epoch, ema_decay, total_epochs, print_freq)\n",
    "    time_elapsed = time.time() - start_time\n",
    "    print('epoch training complete in {: .0f}m {:0f}s'.format(time_elapsed // 60, time_elapsed % 60))\n",
    "    \n",
    "    #validation\n",
    "    start_time = time.time()\n",
    "    print(\"Evaluating the primary model:\")\n",
    "    prec1 = validate(eval_loader, model, global_step, epoch, print_freq = 2)\n",
    "    print(\"Evaluating the EMA model:\")\n",
    "    ema_prec1 = validate(eval_loader, model, global_step, epoch + 1, print_freq = 2)\n",
    "    time_elapsed = time.time() - start_time\n",
    "    print('epoch validation complete in {: .0f}m {:0f}s'.format(time_elapsed // 60, time_elapsed % 60))\n",
    "    \n",
    "    # best precision updation\n",
    "    is_best = ema_prec1 > best_prec1\n",
    "    best_prec1 = max(ema_prec1, best_prec1)\n",
    "    \n",
    "    # saving a checkpoint of the model after the epoch\n",
    "    save_checkpoint({\n",
    "        'epoch': epoch + 1,\n",
    "        'global_step': global_step,\n",
    "        'arch': arch,\n",
    "        'state_dict': model.state_dict(),\n",
    "        'ema_state_dict': ema_model.state_dict(),\n",
    "        'best_prec1': best_prec1,\n",
    "        'optimizer' : optimizer.state_dict(),\n",
    "    }, is_best, checkpoint_path, epoch + 1)"
   ]
  },
  {
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}