train_resnext_attention.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import os, sys\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "import random\n",
    "import config\n",
    "import torch.optim as optim\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "# NEURAL NETWORK MODULES/LAYERS\n",
    "from dataset import Dataset\n",
    "# METRICS CLASS FOR EVALUATION\n",
    "from metrics import Metrics\n",
    "# CONFIG PARSER\n",
    "from config import parse_args\n",
    "# TRAIN AND TEST HELPER FUNCTIONS\n",
    "from trainer import Trainer\n",
    "import glob\n",
    "from tqdm import tqdm_notebook\n",
    "from PIL import Image\n",
    "import torchvision\n",
    "from torchvision import models\n",
    "from tensorboardX import SummaryWriter\n",
    "from torchsummary import summary\n",
    "#import model_vgg16\n",
    "# import model_vgg19\n",
    "# import model_resnet50\n",
    "# import model_densenet121\n",
    "# import model_inceptionv3\n",
    "import model_resnext50\n",
    "\n",
    "# MAIN BLOCK\n",
    "# def main():\n",
    "# global args\n",
    "# args = parse_args()\n",
    "# global logger\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.2.0'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class config():\n",
    "    def __init__(self):\n",
    "        self.data='./data'\n",
    "        self.save='checkpoints/'\n",
    "        self.expname='resnext50_holistic_attention'\n",
    "        self.expno=0\n",
    "        self.pretrained_model='resnext50'\n",
    "        self.num_classes=16\n",
    "        # training arguments\n",
    "        self.epochs=100000\n",
    "        self.batchsize=24\n",
    "        self.lr=0.0001\n",
    "        self.wd=1e-4\n",
    "        self.sparse=False\n",
    "        self.optim='adam'\n",
    "        # miscellaneous options\n",
    "        self.seed=123\n",
    "        self.cuda = torch.cuda.is_available()\n",
    "        self.device = torch.device(\"cuda\" if self.cuda else \"cpu\")\n",
    "        self.pretrained_holistic=3\n",
    "\n",
    "#         cuda_parser = parser.add_mutually_exclusive_group(required=False)\n",
    "#         cuda', dest='cuda', action='store_true')\n",
    "#         cuda_parser.add_argument('--no-cuda', dest='cuda', action='store_false')\n",
    "#         parser.set_defaults(cuda=True)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "24"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "args = config()\n",
    "writer = SummaryWriter(comment = args.expname)\n",
    "args.batchsize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "torch.manual_seed(args.seed)\n",
    "random.seed(args.seed)\n",
    "\n",
    "if args.cuda:\n",
    "    torch.cuda.manual_seed(args.seed)\n",
    "    torch.backends.cudnn.benchmark = True\n",
    "\n",
    "if not os.path.exists(args.save):\n",
    "    os.makedirs(args.save)\n",
    "\n",
    "def plot_data(x,ax):\n",
    "#     x = x.reshape((size,size))\n",
    "    \n",
    "    ax.imshow(x, cmap='gray')\n",
    "#     if y is not None:\n",
    "#         ax.scatter(y[0::2] , y[1::2] , marker='x', s=10)\n",
    "def plot_images(x,batch_size=100):\n",
    "    fig = plt.figure(figsize=(20,20))\n",
    "    fig.subplots_adjust(left=0, right=1, bottom=0, top=1, hspace=0.01, wspace=0.01)\n",
    "\n",
    "    for i in range(batch_size):\n",
    "        ax = fig.add_subplot(10, 10, i + 1, xticks=[], yticks=[])\n",
    "    #     print(y['classes'][i])\n",
    "    #     print(np.argmax(c[i]))\n",
    "        plot_data(x[i], ax)\n",
    "# plot_images(out)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_optimizer(args, parameters):\n",
    "    if args.optim=='adam':\n",
    "        optimizer   = optim.Adam(parameters, lr=args.lr, weight_decay=args.wd)\n",
    "    elif args.optim=='adagrad':\n",
    "        optimizer   = optim.Adagrad(parameters, lr=args.lr, weight_decay=args.wd)\n",
    "    elif args.optim=='sgd':\n",
    "        optimizer   = optim.SGD(parameters, lr=args.lr, weight_decay=args.wd)\n",
    "    elif args.optim == 'adadelta':\n",
    "        optimizer = optim.Adadelta(parameters, lr=args.lr, weight_decay=args.wd)\n",
    "    return optimizer\n",
    "\n",
    "# initialize model, criterion/loss_function, optimizer\n",
    "pretrained_resnext50 = models.resnext50_32x4d(pretrained=True)\n",
    "\n",
    "# Freeze training for all layers\n",
    "# for child in pretrained_vgg16.children():\n",
    "#     for param in child.parameters():\n",
    "#         param.requires_grad = False\n",
    "start_point = 0\n",
    "\n",
    "if args.pretrained_holistic == 0 and 'multiple_loss' not in args.expname:\n",
    "    model = model_resnext50.DocClassificationHolistic(args, pretrained_resnext50)\n",
    "    optimizer = create_optimizer(args, list(model.parameters()))\n",
    "\n",
    "    if os.path.exists('./checkpoints/{0}.pt'.format(args.expname)):\n",
    "        checkpoint = torch.load('./checkpoints/{0}.pt'.format(args.expname))\n",
    "        model.load_state_dict(checkpoint['model'])\n",
    "        if args.cuda:\n",
    "            model.cuda()\n",
    "        #start_point = checkpoint['epoch']\n",
    "        start_point = 0\n",
    "        optimizer.load_state_dict(checkpoint['optim'])\n",
    "elif args.pretrained_holistic == 0 and 'multiple_loss' in args.expname:\n",
    "    model = model_resnext50.DocClassificationHolistic(args, pretrained_resnext50)\n",
    "    optimizer = create_optimizer(args, list(model.parameters()))\n",
    "\n",
    "    if os.path.exists('./checkpoints/{0}.pt'.format(args.expname)):\n",
    "        checkpoint = torch.load('./checkpoints/{0}.pt'.format(args.expname))\n",
    "        model.load_state_dict(checkpoint['model'])\n",
    "        if args.cuda:\n",
    "            model.cuda()\n",
    "        start_point = checkpoint['epoch']\n",
    "        #start_point = 0\n",
    "        optimizer.load_state_dict(checkpoint['optim'])\n",
    "\n",
    "elif args.pretrained_holistic == 1:\n",
    "    checkpoint_holistic = torch.load('./checkpoints/{0}.pt'.format(args.expname))\n",
    "    #resnext50 = model_resnext50.DocClassificationHolistic(args, pretrained_resnext50)\n",
    "    pretrained_holistic_model = model_resnext50.DocClassificationHolistic(args, pretrained_resnext50)\n",
    "    pretrained_holistic_model.load_state_dict(checkpoint_holistic['model'])\n",
    "\n",
    "    model = model_resnext50.DocClassificationRest(args, pretrained_holistic_model)\n",
    "\n",
    "    parameters = list(pretrained_holistic_model.parameters())\n",
    "\n",
    "    if os.path.exists('./checkpoints/%s.pt'%(args.expname)):\n",
    "        checkpoint_ensemble = torch.load('./checkpoints/%s.pt'%(args.expname))\n",
    "        model.load_state_dict(checkpoint_ensemble['model'])\n",
    "        if args.cuda:\n",
    "            model.cuda()\n",
    "\n",
    "        start_point = checkpoint_ensemble['epoch']\n",
    "        optimizer = create_optimizer(args, parameters)\n",
    "        optimizer.load_state_dict(checkpoint_ensemble['optim'])\n",
    "    else:\n",
    "        optimizer = create_optimizer(args, parameters)\n",
    "        #optimizer.load_state_dict(checkpoint_holistic['optim'])elif\n",
    "        \n",
    "elif args.pretrained_holistic == 3:\n",
    "    model = model_resnext50.DocClassification3DAttention(args, pretrained_resnext50)\n",
    "    optimizer = create_optimizer(args, list(model.parameters()))\n",
    "\n",
    "    if os.path.exists('./checkpoints/{0}.pt'.format(args.expname)):\n",
    "        checkpoint = torch.load('./checkpoints/{0}.pt'.format(args.expname))\n",
    "        model.load_state_dict(checkpoint['model'])\n",
    "        if args.cuda:\n",
    "            model.cuda()\n",
    "        #start_point = checkpoint['epoch']\n",
    "        start_point = 0\n",
    "        optimizer.load_state_dict(checkpoint['optim'])\n",
    "criterion = nn.CrossEntropyLoss(reduction='mean')\n",
    "\n",
    "\n",
    "if args.cuda:\n",
    "    model.cuda(), criterion.cuda()\n",
    "\n",
    "\n",
    "metrics = Metrics(args.num_classes)\n",
    "\n",
    "# create trainer object for training and testing\n",
    "trainer = Trainer(args, model, criterion, optimizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [-1, 64, 188, 250]           9,408\n",
      "       BatchNorm2d-2         [-1, 64, 188, 250]             128\n",
      "              ReLU-3         [-1, 64, 188, 250]               0\n",
      "         MaxPool2d-4          [-1, 64, 94, 125]               0\n",
      "            Conv2d-5         [-1, 128, 94, 125]           8,192\n",
      "       BatchNorm2d-6         [-1, 128, 94, 125]             256\n",
      "              ReLU-7         [-1, 128, 94, 125]               0\n",
      "            Conv2d-8         [-1, 128, 94, 125]           4,608\n",
      "       BatchNorm2d-9         [-1, 128, 94, 125]             256\n",
      "             ReLU-10         [-1, 128, 94, 125]               0\n",
      "           Conv2d-11         [-1, 256, 94, 125]          32,768\n",
      "      BatchNorm2d-12         [-1, 256, 94, 125]             512\n",
      "           Conv2d-13         [-1, 256, 94, 125]          16,384\n",
      "      BatchNorm2d-14         [-1, 256, 94, 125]             512\n",
      "             ReLU-15         [-1, 256, 94, 125]               0\n",
      "       Bottleneck-16         [-1, 256, 94, 125]               0\n",
      "           Conv2d-17         [-1, 128, 94, 125]          32,768\n",
      "      BatchNorm2d-18         [-1, 128, 94, 125]             256\n",
      "             ReLU-19         [-1, 128, 94, 125]               0\n",
      "           Conv2d-20         [-1, 128, 94, 125]           4,608\n",
      "      BatchNorm2d-21         [-1, 128, 94, 125]             256\n",
      "             ReLU-22         [-1, 128, 94, 125]               0\n",
      "           Conv2d-23         [-1, 256, 94, 125]          32,768\n",
      "      BatchNorm2d-24         [-1, 256, 94, 125]             512\n",
      "             ReLU-25         [-1, 256, 94, 125]               0\n",
      "       Bottleneck-26         [-1, 256, 94, 125]               0\n",
      "           Conv2d-27         [-1, 128, 94, 125]          32,768\n",
      "      BatchNorm2d-28         [-1, 128, 94, 125]             256\n",
      "             ReLU-29         [-1, 128, 94, 125]               0\n",
      "           Conv2d-30         [-1, 128, 94, 125]           4,608\n",
      "      BatchNorm2d-31         [-1, 128, 94, 125]             256\n",
      "             ReLU-32         [-1, 128, 94, 125]               0\n",
      "           Conv2d-33         [-1, 256, 94, 125]          32,768\n",
      "      BatchNorm2d-34         [-1, 256, 94, 125]             512\n",
      "             ReLU-35         [-1, 256, 94, 125]               0\n",
      "       Bottleneck-36         [-1, 256, 94, 125]               0\n",
      "           Conv2d-37         [-1, 256, 94, 125]          65,536\n",
      "      BatchNorm2d-38         [-1, 256, 94, 125]             512\n",
      "             ReLU-39         [-1, 256, 94, 125]               0\n",
      "           Conv2d-40          [-1, 256, 47, 63]          18,432\n",
      "      BatchNorm2d-41          [-1, 256, 47, 63]             512\n",
      "             ReLU-42          [-1, 256, 47, 63]               0\n",
      "           Conv2d-43          [-1, 512, 47, 63]         131,072\n",
      "      BatchNorm2d-44          [-1, 512, 47, 63]           1,024\n",
      "           Conv2d-45          [-1, 512, 47, 63]         131,072\n",
      "      BatchNorm2d-46          [-1, 512, 47, 63]           1,024\n",
      "             ReLU-47          [-1, 512, 47, 63]               0\n",
      "       Bottleneck-48          [-1, 512, 47, 63]               0\n",
      "           Conv2d-49          [-1, 256, 47, 63]         131,072\n",
      "      BatchNorm2d-50          [-1, 256, 47, 63]             512\n",
      "             ReLU-51          [-1, 256, 47, 63]               0\n",
      "           Conv2d-52          [-1, 256, 47, 63]          18,432\n",
      "      BatchNorm2d-53          [-1, 256, 47, 63]             512\n",
      "             ReLU-54          [-1, 256, 47, 63]               0\n",
      "           Conv2d-55          [-1, 512, 47, 63]         131,072\n",
      "      BatchNorm2d-56          [-1, 512, 47, 63]           1,024\n",
      "             ReLU-57          [-1, 512, 47, 63]               0\n",
      "       Bottleneck-58          [-1, 512, 47, 63]               0\n",
      "           Conv2d-59          [-1, 256, 47, 63]         131,072\n",
      "      BatchNorm2d-60          [-1, 256, 47, 63]             512\n",
      "             ReLU-61          [-1, 256, 47, 63]               0\n",
      "           Conv2d-62          [-1, 256, 47, 63]          18,432\n",
      "      BatchNorm2d-63          [-1, 256, 47, 63]             512\n",
      "             ReLU-64          [-1, 256, 47, 63]               0\n",
      "           Conv2d-65          [-1, 512, 47, 63]         131,072\n",
      "      BatchNorm2d-66          [-1, 512, 47, 63]           1,024\n",
      "             ReLU-67          [-1, 512, 47, 63]               0\n",
      "       Bottleneck-68          [-1, 512, 47, 63]               0\n",
      "           Conv2d-69          [-1, 256, 47, 63]         131,072\n",
      "      BatchNorm2d-70          [-1, 256, 47, 63]             512\n",
      "             ReLU-71          [-1, 256, 47, 63]               0\n",
      "           Conv2d-72          [-1, 256, 47, 63]          18,432\n",
      "      BatchNorm2d-73          [-1, 256, 47, 63]             512\n",
      "             ReLU-74          [-1, 256, 47, 63]               0\n",
      "           Conv2d-75          [-1, 512, 47, 63]         131,072\n",
      "      BatchNorm2d-76          [-1, 512, 47, 63]           1,024\n",
      "             ReLU-77          [-1, 512, 47, 63]               0\n",
      "       Bottleneck-78          [-1, 512, 47, 63]               0\n",
      "           Conv2d-79          [-1, 512, 47, 63]         262,144\n",
      "      BatchNorm2d-80          [-1, 512, 47, 63]           1,024\n",
      "             ReLU-81          [-1, 512, 47, 63]               0\n",
      "           Conv2d-82          [-1, 512, 24, 32]          73,728\n",
      "      BatchNorm2d-83          [-1, 512, 24, 32]           1,024\n",
      "             ReLU-84          [-1, 512, 24, 32]               0\n",
      "           Conv2d-85         [-1, 1024, 24, 32]         524,288\n",
      "      BatchNorm2d-86         [-1, 1024, 24, 32]           2,048\n",
      "           Conv2d-87         [-1, 1024, 24, 32]         524,288\n",
      "      BatchNorm2d-88         [-1, 1024, 24, 32]           2,048\n",
      "             ReLU-89         [-1, 1024, 24, 32]               0\n",
      "       Bottleneck-90         [-1, 1024, 24, 32]               0\n",
      "           Conv2d-91          [-1, 512, 24, 32]         524,288\n",
      "      BatchNorm2d-92          [-1, 512, 24, 32]           1,024\n",
      "             ReLU-93          [-1, 512, 24, 32]               0\n",
      "           Conv2d-94          [-1, 512, 24, 32]          73,728\n",
      "      BatchNorm2d-95          [-1, 512, 24, 32]           1,024\n",
      "             ReLU-96          [-1, 512, 24, 32]               0\n",
      "           Conv2d-97         [-1, 1024, 24, 32]         524,288\n",
      "      BatchNorm2d-98         [-1, 1024, 24, 32]           2,048\n",
      "             ReLU-99         [-1, 1024, 24, 32]               0\n",
      "      Bottleneck-100         [-1, 1024, 24, 32]               0\n",
      "          Conv2d-101          [-1, 512, 24, 32]         524,288\n",
      "     BatchNorm2d-102          [-1, 512, 24, 32]           1,024\n",
      "            ReLU-103          [-1, 512, 24, 32]               0\n",
      "          Conv2d-104          [-1, 512, 24, 32]          73,728\n",
      "     BatchNorm2d-105          [-1, 512, 24, 32]           1,024\n",
      "            ReLU-106          [-1, 512, 24, 32]               0\n",
      "          Conv2d-107         [-1, 1024, 24, 32]         524,288\n",
      "     BatchNorm2d-108         [-1, 1024, 24, 32]           2,048\n",
      "            ReLU-109         [-1, 1024, 24, 32]               0\n",
      "      Bottleneck-110         [-1, 1024, 24, 32]               0\n",
      "          Conv2d-111          [-1, 512, 24, 32]         524,288\n",
      "     BatchNorm2d-112          [-1, 512, 24, 32]           1,024\n",
      "            ReLU-113          [-1, 512, 24, 32]               0\n",
      "          Conv2d-114          [-1, 512, 24, 32]          73,728\n",
      "     BatchNorm2d-115          [-1, 512, 24, 32]           1,024\n",
      "            ReLU-116          [-1, 512, 24, 32]               0\n",
      "          Conv2d-117         [-1, 1024, 24, 32]         524,288\n",
      "     BatchNorm2d-118         [-1, 1024, 24, 32]           2,048\n",
      "            ReLU-119         [-1, 1024, 24, 32]               0\n",
      "      Bottleneck-120         [-1, 1024, 24, 32]               0\n",
      "          Conv2d-121          [-1, 512, 24, 32]         524,288\n",
      "     BatchNorm2d-122          [-1, 512, 24, 32]           1,024\n",
      "            ReLU-123          [-1, 512, 24, 32]               0\n",
      "          Conv2d-124          [-1, 512, 24, 32]          73,728\n",
      "     BatchNorm2d-125          [-1, 512, 24, 32]           1,024\n",
      "            ReLU-126          [-1, 512, 24, 32]               0\n",
      "          Conv2d-127         [-1, 1024, 24, 32]         524,288\n",
      "     BatchNorm2d-128         [-1, 1024, 24, 32]           2,048\n",
      "            ReLU-129         [-1, 1024, 24, 32]               0\n",
      "      Bottleneck-130         [-1, 1024, 24, 32]               0\n",
      "          Conv2d-131          [-1, 512, 24, 32]         524,288\n",
      "     BatchNorm2d-132          [-1, 512, 24, 32]           1,024\n",
      "            ReLU-133          [-1, 512, 24, 32]               0\n",
      "          Conv2d-134          [-1, 512, 24, 32]          73,728\n",
      "     BatchNorm2d-135          [-1, 512, 24, 32]           1,024\n",
      "            ReLU-136          [-1, 512, 24, 32]               0\n",
      "          Conv2d-137         [-1, 1024, 24, 32]         524,288\n",
      "     BatchNorm2d-138         [-1, 1024, 24, 32]           2,048\n",
      "            ReLU-139         [-1, 1024, 24, 32]               0\n",
      "      Bottleneck-140         [-1, 1024, 24, 32]               0\n",
      "          Conv2d-141         [-1, 1024, 24, 32]       1,048,576\n",
      "     BatchNorm2d-142         [-1, 1024, 24, 32]           2,048\n",
      "            ReLU-143         [-1, 1024, 24, 32]               0\n",
      "          Conv2d-144         [-1, 1024, 12, 16]         294,912\n",
      "     BatchNorm2d-145         [-1, 1024, 12, 16]           2,048\n",
      "            ReLU-146         [-1, 1024, 12, 16]               0\n",
      "          Conv2d-147         [-1, 2048, 12, 16]       2,097,152\n",
      "     BatchNorm2d-148         [-1, 2048, 12, 16]           4,096\n",
      "          Conv2d-149         [-1, 2048, 12, 16]       2,097,152\n",
      "     BatchNorm2d-150         [-1, 2048, 12, 16]           4,096\n",
      "            ReLU-151         [-1, 2048, 12, 16]               0\n",
      "      Bottleneck-152         [-1, 2048, 12, 16]               0\n",
      "          Conv2d-153         [-1, 1024, 12, 16]       2,097,152\n",
      "     BatchNorm2d-154         [-1, 1024, 12, 16]           2,048\n",
      "            ReLU-155         [-1, 1024, 12, 16]               0\n",
      "          Conv2d-156         [-1, 1024, 12, 16]         294,912\n",
      "     BatchNorm2d-157         [-1, 1024, 12, 16]           2,048\n",
      "            ReLU-158         [-1, 1024, 12, 16]               0\n",
      "          Conv2d-159         [-1, 2048, 12, 16]       2,097,152\n",
      "     BatchNorm2d-160         [-1, 2048, 12, 16]           4,096\n",
      "            ReLU-161         [-1, 2048, 12, 16]               0\n",
      "      Bottleneck-162         [-1, 2048, 12, 16]               0\n",
      "          Conv2d-163         [-1, 1024, 12, 16]       2,097,152\n",
      "     BatchNorm2d-164         [-1, 1024, 12, 16]           2,048\n",
      "            ReLU-165         [-1, 1024, 12, 16]               0\n",
      "          Conv2d-166         [-1, 1024, 12, 16]         294,912\n",
      "     BatchNorm2d-167         [-1, 1024, 12, 16]           2,048\n",
      "            ReLU-168         [-1, 1024, 12, 16]               0\n",
      "          Conv2d-169         [-1, 2048, 12, 16]       2,097,152\n",
      "     BatchNorm2d-170         [-1, 2048, 12, 16]           4,096\n",
      "            ReLU-171         [-1, 2048, 12, 16]               0\n",
      "      Bottleneck-172         [-1, 2048, 12, 16]               0\n",
      "          Conv3d-173         [-1, 1, 1, 12, 16]          18,433\n",
      "          Linear-174                  [-1, 256]         524,544\n",
      "          Linear-175                   [-1, 16]           4,112\n",
      "================================================================\n",
      "Total params: 23,526,993\n",
      "Trainable params: 23,526,993\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 2.15\n",
      "Forward/backward pass size (MB): 1373.11\n",
      "Params size (MB): 89.75\n",
      "Estimated Total Size (MB): 1465.00\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "summary(model,(3,375,500))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DocClassification3DAttention(\n",
       "  (pretrained_model): Sequential(\n",
       "    (0): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "    (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): ReLU(inplace=True)\n",
       "    (3): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "    (4): Sequential(\n",
       "      (0): Bottleneck(\n",
       "        (conv1): Conv2d(64, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (downsample): Sequential(\n",
       "          (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (1): Bottleneck(\n",
       "        (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (2): Bottleneck(\n",
       "        (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "    )\n",
       "    (5): Sequential(\n",
       "      (0): Bottleneck(\n",
       "        (conv1): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (downsample): Sequential(\n",
       "          (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (1): Bottleneck(\n",
       "        (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (2): Bottleneck(\n",
       "        (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (3): Bottleneck(\n",
       "        (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "    )\n",
       "    (6): Sequential(\n",
       "      (0): Bottleneck(\n",
       "        (conv1): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (downsample): Sequential(\n",
       "          (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "          (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (1): Bottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (2): Bottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (3): Bottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (4): Bottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (5): Bottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "    )\n",
       "    (7): Sequential(\n",
       "      (0): Bottleneck(\n",
       "        (conv1): Conv2d(1024, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (downsample): Sequential(\n",
       "          (0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "          (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (1): Bottleneck(\n",
       "        (conv1): Conv2d(2048, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (2): Bottleneck(\n",
       "        (conv1): Conv2d(2048, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (V): Conv3d(1, 1, kernel_size=(2048, 3, 3), stride=(2048, 1, 1), padding=(0, 1, 1))\n",
       "  (fc): Linear(in_features=2048, out_features=256, bias=True)\n",
       "  (out): Linear(in_features=256, out_features=16, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model\n",
    "# for child, name in zip(model.children(), model.named_parameters(), model.n):\n",
    "#     for param in child.parameters():\n",
    "#         param.requires_grad = True\n",
    "#         print(name[0],':',param.shape)\n",
    "#     break\n",
    "# for parameter in model.parameters():\n",
    "#     print(parameter)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==> Size of train data   : 80000 \n",
      "==> Size of val data   : 16000 \n",
      "==> Size of test data   : 15999 \n"
     ]
    }
   ],
   "source": [
    "# 'data/labels/train.csv'\n",
    "train_dataset = Dataset('data/stratified/rvl-cdip/labels/train.csv', model_type=args.pretrained_model)\n",
    "dev_dataset = Dataset('data/stratified/rvl-cdip/labels/val.csv', model_type=args.pretrained_model)\n",
    "test_dataset = Dataset('data/stratified/rvl-cdip/labels/test.csv', model_type=args.pretrained_model)\n",
    "\n",
    "print('==> Size of train data   : %d ' % len(train_dataset))\n",
    "print('==> Size of val data   : %d ' % len(dev_dataset))\n",
    "print('==> Size of test data   : %d ' % len(test_dataset))\n",
    "\n",
    "train_idx = list(np.arange(len(train_dataset)))\n",
    "dev_idx = list(np.arange(len(dev_dataset)))\n",
    "test_idx = list(np.arange(len(test_dataset)))\n",
    "\n",
    "best = -float('inf')\n",
    "columns = ['ExpName','ExpNo', 'Epoch', 'Loss','Accuracy']\n",
    "results = []\n",
    "early_stop_count = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# bh,h,f,l,r,lab = train_dataset[np.random.randint(0,5000,16)]\n",
    "# batch_holistic = torch.unsqueeze(bh , dim=1).repeat(1,3,1,1).float().cuda()\n",
    "# # bh = np.array(bh.detach().cpu())\n",
    "# print(bh.shape, batch_holistic.shape)\n",
    "# batch_header = torch.unsqueeze(h, dim=1).repeat(1,3,1,1).float().cuda()\n",
    "# batch_footer = torch.unsqueeze(f, dim=1).repeat(1,3,1,1).float().cuda()\n",
    "# batch_left = torch.unsqueeze(l, dim=1).repeat(1,3,1,1).float().cuda()\n",
    "# batch_right = torch.unsqueeze(r, dim=1).repeat(1,3,1,1).float().cuda()\n",
    "\n",
    "# all_data = [batch_holistic]\n",
    "# writer.add_graph(model, input_to_model=all_data)\n",
    "# # writer.flush()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from matplotlib import pyplot as plt\n",
    "# def plot_data(x,ax):\n",
    "# #     x = x.reshape((size,size))\n",
    "    \n",
    "#     ax.imshow(x, cmap='gray')\n",
    "# #     if y is not None:\n",
    "# #         ax.scatter(y[0::2] , y[1::2] , marker='x', s=10)\n",
    "# def plot_images(x,batch_size=64):\n",
    "#     fig = plt.figure(figsize=(50,50))\n",
    "#     fig.subplots_adjust(left=0, right=1, bottom=0, top=1, hspace=0.01, wspace=0.01)\n",
    "#     for i in range(batch_size):\n",
    "#         ax = fig.add_subplot(8, 8, i + 1, xticks=[], yticks=[])\n",
    "#     #     print(y['classes'][i])\n",
    "#     #     print(np.argmax(c[i]))\n",
    "#         plot_data(x[i], ax)\n",
    "# plot_images(bh)\n",
    "# print(l)\n",
    "# # display(plot_images(np.squeeze(seg[0],axis=-1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fc644679886c41b7ac28d3ef8b6fe50b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='training', max=100000, style=ProgressStyle(description_width=…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "86cf4711d0994bc0a7fea9cda5729f59",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Testing', max=667, style=ProgressStyle(description_width='ini…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "for epoch in tqdm_notebook(range(start_point, args.epochs), desc= 'training', leave=True):\n",
    "\n",
    "    train_loss = 0.0\n",
    "    dev_loss = 0.0\n",
    "    test_loss = 0.0\n",
    "\n",
    "    train_predictions = []\n",
    "    train_labels = []\n",
    "\n",
    "    dev_predictions = []\n",
    "    dev_labels = []\n",
    "\n",
    "    test_predictions = []\n",
    "    test_labels = []\n",
    "\n",
    "    \n",
    "    batch_test_idxs = []\n",
    "    counter = 0\n",
    "\n",
    "    while counter<len(test_dataset):\n",
    "        batch_test_idx = np.array(list(np.arange(counter, counter+args.batchsize)))\n",
    "        batch_test_idxs.append(batch_test_idx)\n",
    "        counter+=args.batchsize\n",
    "    batch_test_idxs\n",
    "\n",
    "    test_acc = []\n",
    "    test_predictions = []\n",
    "    test_labels = []\n",
    "    \n",
    "\n",
    "    batch_train_data = list(np.random.randint(0, len(train_idx), args.batchsize))\n",
    "    batch_dev_data = list(np.random.randint(0, len(dev_idx), args.batchsize))\n",
    "    batch_test_data = list(np.random.randint(0, len(test_idx), args.batchsize))\n",
    "    \n",
    "#     for batch in tqdm(batch_train_data, desc='Training batches..'):\n",
    "    # ----------------- Train ---------------------\n",
    "    train_batch_holistic, \\\n",
    "            train_batch_header, \\\n",
    "            train_batch_footer, \\\n",
    "            train_batch_left_body, \\\n",
    "            train_batch_right_body, \\\n",
    "            train_batch_labels = train_dataset[batch_train_data]\n",
    "#     print(train_batch_holistic.shape, \\\n",
    "#             train_batch_header.shape, \\\n",
    "#             train_batch_footer.shape, \\\n",
    "#             train_batch_left_body.shape, \\\n",
    "#             train_batch_right_body.shape, \\\n",
    "#             train_batch_labels.shape)\n",
    "    if args.pretrained_holistic == 0 and args.expname == 'resnext50_holistic':\n",
    "        train_batch_loss,train_batch_predictions,train_batch_labels = trainer.train_holistic(train_batch_holistic, train_batch_labels)\n",
    "    elif args.pretrained_holistic == 3:\n",
    "        train_batch_loss,train_batch_predictions,train_batch_labels, train_attn_map = trainer.train_3D_attention(train_batch_holistic, train_batch_labels)    \n",
    "    elif args.pretrained_holistic == 0 and args.expname != 'resnext50_holistic':\n",
    "        train_batch_loss,train_batch_predictions,train_batch_labels = trainer.train_rest_multiple_loss(train_batch_holistic, \\\n",
    "                                                                                    train_batch_header, \\\n",
    "                                                                                    train_batch_footer, \\\n",
    "                                                                                    train_batch_left_body, \\\n",
    "                                                                                    train_batch_right_body, \\\n",
    "                                                                                    train_batch_labels)\n",
    "    elif args.pretrained_holistic == 1:\n",
    "        train_batch_loss,train_batch_predictions,train_batch_labels = trainer.train_rest(train_batch_holistic, \\\n",
    "                                train_batch_header, \\\n",
    "                                train_batch_footer, \\\n",
    "                                train_batch_left_body, \\\n",
    "                                train_batch_right_body, \\\n",
    "                                train_batch_labels)\n",
    "\n",
    "\n",
    "    train_accuracy = metrics.accuracy(train_batch_predictions, train_batch_labels)\n",
    "    # ----------------- Dev ---------------------\n",
    "#     for batch in tqdm(batch_dev_data, desc='Dev batches..'):\n",
    "    dev_batch_holistic, \\\n",
    "            dev_batch_header, \\\n",
    "            dev_batch_footer, \\\n",
    "            dev_batch_left_body, \\\n",
    "            dev_batch_right_body, \\\n",
    "            dev_batch_labels = dev_dataset[batch_dev_data]\n",
    "\n",
    "    if args.pretrained_holistic == 0 and args.expname == 'resnext50_holistic':\n",
    "        dev_batch_loss, dev_batch_predictions, dev_batch_labels = trainer.test_holistic(dev_batch_holistic, dev_batch_labels)\n",
    "    elif args.pretrained_holistic == 3:\n",
    "        dev_batch_loss, dev_batch_predictions, dev_batch_labels, dev_attn_map = trainer.test_3D_attention(dev_batch_holistic, dev_batch_labels)    \n",
    "    \n",
    "    elif args.pretrained_holistic == 0 and args.expname != 'resnext50_holistic':\n",
    "        dev_batch_loss, dev_batch_predictions, dev_batch_labels = trainer.test_rest_multiple_loss(dev_batch_holistic, \\\n",
    "                                                                                dev_batch_header, \\\n",
    "                                                                                dev_batch_footer, \\\n",
    "                                                                                dev_batch_left_body, \\\n",
    "                                                                                dev_batch_right_body, \\\n",
    "                                                                                dev_batch_labels)\n",
    "    elif args.pretrained_holistic == 1:\n",
    "        dev_batch_loss, dev_batch_predictions, dev_batch_labels = trainer.test_rest(dev_batch_holistic, \\\n",
    "                                dev_batch_header, \\\n",
    "                                dev_batch_footer, \\\n",
    "                                dev_batch_left_body, \\\n",
    "                                dev_batch_right_body, \\\n",
    "                                dev_batch_labels)\n",
    "#         dev_predictions.append(dev_batch_predictions)\n",
    "#         dev_labels.append(dev_batch_labels)\n",
    "#         dev_loss = dev_loss + dev_batch_loss\n",
    "\n",
    "    dev_accuracy = metrics.accuracy(dev_batch_predictions, dev_batch_labels)\n",
    "    # ----------------- Test ---------------------\n",
    "    for batch_test_idx in tqdm_notebook(batch_test_idxs, desc='Testing', leave=False ):\n",
    "#         print(np.round(batch_test_idx[-1]*100/16000,2), end='\\r')\n",
    "        try:\n",
    "            test_batch_holistic, \\\n",
    "                    test_batch_header, \\\n",
    "                    test_batch_footer, \\\n",
    "                    test_batch_left_body, \\\n",
    "                    test_batch_right_body, \\\n",
    "                    test_batch_labels = test_dataset[batch_test_idx]\n",
    "\n",
    "            if args.pretrained_holistic == 0 and 'multiple_loss' not in args.expname: # singleloss\n",
    "                test_batch_loss, test_batch_predictions, test_batch_labels = trainer.test_holistic(test_batch_header, test_batch_labels)\n",
    "            elif args.pretrained_holistic == 3:\n",
    "                test_batch_loss, test_batch_predictions, test_batch_labels, test_attn_map = trainer.test_3D_attention(test_batch_header, test_batch_labels)\n",
    "            elif args.pretrained_holistic == 0 and 'multiple_loss' in args.expname: # multipleloss\n",
    "                test_batch_loss, test_batch_predictions, test_batch_labels = trainer.test_rest_multiple_loss(test_batch_holistic, \\\n",
    "                                                                                        test_batch_header, \\\n",
    "                                                                                        test_batch_footer, \\\n",
    "                                                                                        test_batch_left_body, \\\n",
    "                                                                                        test_batch_right_body, \\\n",
    "                                                                                        test_batch_labels)\n",
    "            elif args.pretrained_holistic == 1:\n",
    "                test_batch_loss, test_batch_predictions, test_batch_labels = trainer.test_rest(test_batch_holistic, \\\n",
    "                                        test_batch_header, \\\n",
    "                                        test_batch_footer, \\\n",
    "                                        test_batch_left_body, \\\n",
    "                                        test_batch_right_body, \\\n",
    "                                        test_batch_labels)\n",
    "\n",
    "            elif args.pretrained_holistic == 2:\n",
    "    #             print('here')\n",
    "                test_batch_loss, test_batch_predictions, test_batch_labels = trainer.test_ensemble(test_batch_holistic, \\\n",
    "                                        test_batch_header, \\\n",
    "                                        test_batch_footer, \\\n",
    "                                        test_batch_left_body, \\\n",
    "                                        test_batch_right_body, \\\n",
    "                                        test_batch_labels)\n",
    "        #         test_predictions.append(test_batch_predictions)\n",
    "        #         test_labels.append(test_batch_labels)\n",
    "        #         test_loss = test_loss + test_batch_loss\n",
    "\n",
    "            test_accuracy = metrics.accuracy(test_batch_predictions, test_batch_labels)\n",
    "        #     print(test_accuracy)\n",
    "            test_acc.append(test_accuracy)\n",
    "        except Exception as ex:\n",
    "            pass #print(ex)\n",
    "\n",
    "\n",
    "    test_accuracy = np.mean(test_acc) # metrics.accuracy(test_batch_predictions, test_batch_labels)\n",
    "    \n",
    "    # ----------------- Scalars ---------------------\n",
    "    writer.add_scalar('Loss/train',train_batch_loss.detach().cpu(), epoch)\n",
    "    writer.add_scalar('Loss/test', test_batch_loss.detach().cpu(), epoch)\n",
    "    writer.add_scalar('Loss/val', dev_batch_loss.detach().cpu(), epoch)\n",
    "    writer.add_scalar('Accuracy/val', dev_accuracy, epoch)\n",
    "    writer.add_scalar('Accuracy/train', train_accuracy, epoch)\n",
    "    writer.add_scalar('Accuracy/test', test_accuracy, epoch)\n",
    "\n",
    "    if best < test_accuracy:\n",
    "        best = test_accuracy\n",
    "        checkpoint = {'model': trainer.model.state_dict(), 'optim': trainer.optimizer.state_dict(),\n",
    "                      'loss': test_batch_loss, 'accuracy': test_accuracy,\n",
    "                      'args': args, 'epoch': epoch }\n",
    "#         logger.debug('==> New optimum found, checkpointing everything now...')\n",
    "        torch.save(checkpoint, '%s.pt' % os.path.join(args.save, args.expname))\n",
    "        #np.savetxt(\"test_pred.csv\", test_pred.numpy(), delimiter=\",\")\n",
    "        early_stop_count = 0\n",
    "    else:\n",
    "        early_stop_count = early_stop_count + 1\n",
    "\n",
    "        if early_stop_count >= 20:\n",
    "            print('Early Stopping')\n",
    "            print('Epoch {:7d}, Train Loss: {:.2f}, Train Accuracy: {:.2f},' \\\n",
    "          'Dev Loss: {:.2f}, Dev Accuracy: {:.2f}, Test Loss: {:.2f},' \\\n",
    "          'Test Accuracy: {:.2f}'.format(epoch + 1 \\\n",
    "                                     , train_batch_loss.detach().cpu() \\\n",
    "                                     , train_accuracy \\\n",
    "                                     , dev_batch_loss.detach().cpu() \\\n",
    "                                     , dev_accuracy \\\n",
    "                                     , test_batch_loss.detach().cpu() \\\n",
    "                                     , test_accuracy))\n",
    "    \n",
    "            break\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# bh = train_dataset[np.random.randint(0,10000,64)][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# bh = train_dataset[np.random.randint(0,10000,64)][0]\n",
    "# batch_holistic = torch.unsqueeze(bh, dim=1).repeat(1,3,1,1).float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os,numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# l = np.array(sorted(os.listdir('data/test/holistic/')))\n",
    "# _,cts = np.unique(l,return_counts=True)\n",
    "# cts[cts>1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pd.read_csv('data/labels/test.csv',sep=' ')\n",
    "# np.unique(l,return_counts=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import numpy as np\n",
    "# from PIL import Image\n",
    "# # Image.open('data/test/holistic/15999.jpg', mode='r')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# len(os.listdir('data/test/holistic/'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# df_labels = pd.read_csv('data/stratified/rvl-cdip/labels/train.csv',names=['fn','l'],sep=' ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filenames = df_labels[df_labels['l']==12]['fn'].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# idx = np.random.randint(0,5000,64)\n",
    "# images = []\n",
    "# for f in filenames[idx]:\n",
    "#     images.append(np.array(Image.open(os.path.join('data/stratified/rvl-cdip/images/',f))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot_images(images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (pt-gpu)",
   "language": "python",
   "name": "pytorch-gpu"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}