train_resnet_375x500.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cd /research4/projects/document_classification/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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",
    "\n",
    "#import model_vgg16\n",
    "# import model_vgg19\n",
    "import model_resnet50\n",
    "# import model_densenet121\n",
    "# import model_inceptionv3\n",
    "\n",
    "# MAIN BLOCK\n",
    "# def main():\n",
    "# global args\n",
    "# args = parse_args()\n",
    "# global logger\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.2.0'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class config():\n",
    "    def __init__(self):\n",
    "        self.data='./data'\n",
    "        self.save='checkpoints/'\n",
    "        self.expname='resnet50_holistic'\n",
    "        self.expno=0\n",
    "        self.pretrained_model='resnet50'\n",
    "        self.num_classes=16\n",
    "        self.pretrained_holistic=0\n",
    "        # training arguments\n",
    "        self.epochs=100000\n",
    "        self.batchsize=30\n",
    "        self.lr=0.0001\n",
    "        self.wd=1e-4\n",
    "        self.sparse=False\n",
    "        self.optim='adam'\n",
    "        self.input_size = '_375x500'\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",
    "#         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)        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "30"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "args = config()\n",
    "writer = SummaryWriter(comment = '{0}{1}'.format(args.expname,args.input_size))\n",
    "args.batchsize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "here\n",
      "model exists\n",
      "start_point: 1\n"
     ]
    }
   ],
   "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_resnet50 = models.resnet50(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_resnet50.DocClassificationHolistic(args, pretrained_resnet50)\n",
    "    optimizer = create_optimizer(args, list(model.parameters()))\n",
    "    print('here')\n",
    "    if os.path.exists('./checkpoints/{0}{1}.pt'.format(args.expname,args.input_size)):\n",
    "        print('model exists')\n",
    "        checkpoint = torch.load('./checkpoints/{0}{1}.pt'.format(args.expname,args.input_size))\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_resnet50.DocClassificationHolistic(args, pretrained_resnet50)\n",
    "    optimizer = create_optimizer(args, list(model.parameters()))\n",
    "\n",
    "    if os.path.exists('./checkpoints/{0}{1}.pt'.format(args.expname,args.input_size)):\n",
    "        checkpoint = torch.load('./checkpoints/{0}{1}.pt'.format(args.expname,args.input_size))\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}{1}.pt'.format(args.expname,args.input_size))\n",
    "    #resnet50 = model_resnet50.DocClassificationHolistic(args, pretrained_resnet50)\n",
    "    pretrained_holistic_model = model_resnet50.DocClassificationHolistic(args, pretrained_resnet50)\n",
    "    pretrained_holistic_model.load_state_dict(checkpoint_holistic['model'])\n",
    "\n",
    "    model = model_resnet50.DocClassificationRest(args, pretrained_holistic_model)\n",
    "\n",
    "    parameters = list(pretrained_holistic_model.parameters())\n",
    "\n",
    "    if os.path.exists('./checkpoints/{0}{1}.pt'.format(args.expname, args.input_size)):\n",
    "        checkpoint_ensemble = torch.load('./checkpoints/{0}{1}.pt'.format(args.expname, args.input_size))\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'])\n",
    "\n",
    "\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)\n",
    "print('start_point:',start_point)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # for child in pretrained_vgg16.children():\n",
    "# #     for param in child.parameters():\n",
    "# #         param.requires_grad = True\n",
    "# #         print(param)\n",
    "# for parameter in model.parameters():\n",
    "#     print(parameter)\n",
    "# # model\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\n",
    "\n",
    "# ---------------------------------------------------------------- #\n",
    "# train_dir = glob.glob(os.path.join(args.data,'train/holistic/*.jpg'))\n",
    "# dev_dir = glob.glob(os.path.join(args.data,'val/holistic/*.jpg'))\n",
    "# test_dir = glob.glob(os.path.join(args.data,'test/holistic/*.jpg'))\n",
    "\n",
    "# print(len(train_dir), len(dev_dir), len(test_dir))\n",
    "\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": [
    {
     "data": {
      "text/plain": [
       "torch.Size([64, 500, 375])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bh,h,f,l,r,lab = train_dataset[np.random.randint(55000,60000,64)]\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,batch_header,batch_footer,batch_left,batch_right]\n",
    "# writer.add_graph(model, input_to_model=all_data)\n",
    "# # writer.flush()\n",
    "bh.shape"
   ]
  },
  {
   "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": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8fa92a744f904231985936e77bc0ecbe",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='training', max=99999, style=ProgressStyle(description_width='…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-12-ee721a2c57f6>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     36\u001b[0m \u001b[0;31m#             train_batch_labels.shape)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     37\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpretrained_holistic\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m0\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0;34m'multiple_loss'\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32min\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexpname\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;31m# singleloss\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 38\u001b[0;31m         \u001b[0mtrain_batch_loss\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mtrain_batch_predictions\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mtrain_batch_labels\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain_holistic\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_batch_holistic\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain_batch_labels\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     39\u001b[0m     \u001b[0;32melif\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpretrained_holistic\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m0\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0;34m'multiple_loss'\u001b[0m \u001b[0;32min\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexpname\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;31m# multipleloss\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     40\u001b[0m         train_batch_loss,train_batch_predictions,train_batch_labels = trainer.train_rest_multiple_loss(train_batch_holistic, \\\n",
      "\u001b[0;32m~/AdvancedProjects/Document-Classification/trainer.py\u001b[0m in \u001b[0;36mtrain_holistic\u001b[0;34m(self, batch_holistic, labels)\u001b[0m\n\u001b[1;32m     34\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     35\u001b[0m         \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 36\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     37\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\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     38\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/pytorch-gpu/lib/python3.7/site-packages/torch/optim/adam.py\u001b[0m in \u001b[0;36mstep\u001b[0;34m(self, closure)\u001b[0m\n\u001b[1;32m     92\u001b[0m                 \u001b[0;31m# Decay the first and second moment running average coefficient\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     93\u001b[0m                 \u001b[0mexp_avg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmul_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbeta1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mbeta1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 94\u001b[0;31m                 \u001b[0mexp_avg_sq\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmul_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbeta2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maddcmul_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mbeta2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     95\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0mamsgrad\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     96\u001b[0m                     \u001b[0;31m# Maintains the maximum of all 2nd moment running avg. till now\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# early_stop_count = 0\n",
    "# start_point = 100000 - 65568 - 1873\n",
    "for epoch in tqdm_notebook(range(start_point, args.epochs), desc= 'training'):\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",
    "\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",
    "    # ----------------- Train ---------------------\n",
    "#     for batch in tqdm_notebook(batch_train_data, desc='Training batches..'):\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 'multiple_loss' not in args.expname: # singleloss\n",
    "        train_batch_loss,train_batch_predictions,train_batch_labels = trainer.train_holistic(train_batch_holistic, train_batch_labels)\n",
    "    elif args.pretrained_holistic == 0 and 'multiple_loss' in args.expname: # multipleloss\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_notebook(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 'multiple_loss' not in args.expname: # singleloss\n",
    "        dev_batch_loss, dev_batch_predictions, dev_batch_labels = trainer.test_holistic(dev_batch_holistic, dev_batch_labels)\n",
    "    elif args.pretrained_holistic == 0 and 'multiple_loss' in args.expname: # multipleloss\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 in tqdm_notebook(batch_test_data, desc='Test batches..'):\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_data]\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_holistic, 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",
    "#         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",
    "    \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_batch_loss:\n",
    "        best = test_batch_loss\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, '{0}{1}.pt'.format(os.path.join(args.save, args.expname),args.input_size))\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 >= 1250*10:\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": [
    "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",
    "# ----------------- Test ---------------------\n",
    "for batch_test_idx in batch_test_idxs:\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",
    "#             print('here')\n",
    "            test_batch_loss, test_batch_predictions, test_batch_labels = trainer.test_holistic(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",
    "        print(ex)\n",
    "        \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'Accuracy of {0}: {1} %'.format(args.expname, np.round(np.mean(test_acc)*100,2))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_acc"
   ]
  },
  {
   "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
}