MIL_MNIST/end_to_end_mnist_MIL.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/linqisheng/anaconda3/lib/python3.7/site-packages/tqdm/autonotebook.py:17: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n",
      "  \" (e.g. in jupyter console)\", TqdmExperimentalWarning)\n",
      "/home/linqisheng/anaconda3/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:526: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/home/linqisheng/anaconda3/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:527: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/home/linqisheng/anaconda3/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:528: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/home/linqisheng/anaconda3/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:529: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/home/linqisheng/anaconda3/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:530: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/home/linqisheng/anaconda3/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:535: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import datetime\n",
    "import copy\n",
    "import re\n",
    "import yaml\n",
    "import uuid\n",
    "import warnings\n",
    "import time\n",
    "import inspect\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from functools import partial, reduce\n",
    "from random import shuffle\n",
    "import random\n",
    "\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data.dataset import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision.models import resnet\n",
    "from torchvision.transforms import Compose, ToTensor, Normalize, Resize\n",
    "from torchvision.models.resnet import ResNet, BasicBlock\n",
    "from torchvision.datasets import MNIST\n",
    "import tensorflow as tf\n",
    "from tqdm.autonotebook import tqdm\n",
    "from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score\n",
    "from sklearn import metrics as mtx\n",
    "from sklearn import model_selection as ms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> # 1. Pretrain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data_loaders(train_batch_size, val_batch_size):\n",
    "    mnist = MNIST(download=True, train=True, root=\".\").train_data.float()\n",
    "    \n",
    "    data_transform = Compose([ Resize((224, 224)),ToTensor(), Normalize((mnist.mean()/255,), (mnist.std()/255,))])\n",
    "\n",
    "    train_loader = DataLoader(MNIST(download=True, root=\".\", transform=data_transform, train=True),\n",
    "                              batch_size=train_batch_size, shuffle=True)\n",
    "    val_loader = DataLoader(MNIST(download=False, root=\".\", transform=data_transform, train=False),\n",
    "                            batch_size=val_batch_size, shuffle=False)\n",
    "    return train_loader, val_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/linqisheng/anaconda3/lib/python3.7/site-packages/torchvision/datasets/mnist.py:53: UserWarning: train_data has been renamed data\n",
      "  warnings.warn(\"train_data has been renamed data\")\n"
     ]
    }
   ],
   "source": [
    "train_batch_size = 256\n",
    "val_batch_size = 256\n",
    "\n",
    "train_loader, valid_loader = get_data_loaders(train_batch_size, val_batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> ### 1.1 define model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MnistResNet(ResNet):\n",
    "    def __init__(self):\n",
    "        super(MnistResNet, self).__init__(BasicBlock, [2, 2, 2, 2], num_classes=10)\n",
    "        self.conv1 = torch.nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return torch.softmax(super(MnistResNet, self).forward(x), dim=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> ### 1.2 helper function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_metric(metric_fn, true_y, pred_y):\n",
    "    # multi class problems need to have averaging method\n",
    "    if \"average\" in inspect.getfullargspec(metric_fn).args:\n",
    "        return metric_fn(true_y, pred_y, average=\"macro\")\n",
    "    else:\n",
    "        return metric_fn(true_y, pred_y)\n",
    "    \n",
    "def print_scores(p, r, f1, a, batch_size):\n",
    "    # just an utility printing function\n",
    "    for name, scores in zip((\"precision\", \"recall\", \"F1\", \"accuracy\"), (p, r, f1, a)):\n",
    "        print(f\"\\t{name.rjust(14, ' ')}: {sum(scores)/batch_size:.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> ### 1.3 train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/linqisheng/anaconda3/lib/python3.7/site-packages/torchvision/datasets/mnist.py:53: UserWarning: train_data has been renamed data\n",
      "  warnings.warn(\"train_data has been renamed data\")\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d923a68a9658484394a06b959625a108",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Loss: ', max=235, style=ProgressStyle(description_width='init…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 1/5, training loss: 1.6542498821907856, validation loss: 1.5697572231292725\n",
      "\t     precision: 0.9335\n",
      "\t        recall: 0.8941\n",
      "\t            F1: 0.8905\n",
      "\t      accuracy: 0.8975\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/linqisheng/anaconda3/lib/python3.7/site-packages/sklearn/metrics/classification.py:1437: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 in labels with no predicted samples.\n",
      "  'precision', 'predicted', average, warn_for)\n",
      "/home/linqisheng/anaconda3/lib/python3.7/site-packages/sklearn/metrics/classification.py:1437: UndefinedMetricWarning: F-score is ill-defined and being set to 0.0 in labels with no predicted samples.\n",
      "  'precision', 'predicted', average, warn_for)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "97d1fcc7cb0a4d219320bee253c54da9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Loss: ', max=235, style=ProgressStyle(description_width='init…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 2/5, training loss: 1.4776017858626995, validation loss: 1.5797227621078491\n",
      "\t     precision: 0.9231\n",
      "\t        recall: 0.8847\n",
      "\t            F1: 0.8794\n",
      "\t      accuracy: 0.8869\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7523f4f87c454d1ca2605e48f3ed9427",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Loss: ', max=235, style=ProgressStyle(description_width='init…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 3/5, training loss: 1.4717868257076183, validation loss: 1.4733632802963257\n",
      "\t     precision: 0.9905\n",
      "\t        recall: 0.9904\n",
      "\t            F1: 0.9902\n",
      "\t      accuracy: 0.9903\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a9c8ca9944404ff5b1962f8da867ea3a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Loss: ', max=235, style=ProgressStyle(description_width='init…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 4/5, training loss: 1.4688895022615474, validation loss: 1.4780795574188232\n",
      "\t     precision: 0.9864\n",
      "\t        recall: 0.9866\n",
      "\t            F1: 0.9860\n",
      "\t      accuracy: 0.9865\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "04fd8235e3e84c34bf70b0cedb965124",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Loss: ', max=235, style=ProgressStyle(description_width='init…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 5/5, training loss: 1.4669518562073403, validation loss: 1.4729316234588623\n",
      "\t     precision: 0.9892\n",
      "\t        recall: 0.9890\n",
      "\t            F1: 0.9888\n",
      "\t      accuracy: 0.9891\n",
      "Training time: 618.7201232910156s\n"
     ]
    }
   ],
   "source": [
    "start_ts = time.time()\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "\n",
    "# model:\n",
    "model = MnistResNet().to(device)\n",
    "\n",
    "# params you need to specify:\n",
    "epochs = 5\n",
    "train_loader, val_loader = get_data_loaders(train_batch_size, val_batch_size)\n",
    "loss_function = nn.CrossEntropyLoss() # your loss function, cross entropy works well for multi-class problems\n",
    "\n",
    "# optimizer, I've used Adadelta, as it wokrs well without any magic numbers\n",
    "optimizer = optim.Adadelta(model.parameters())\n",
    "\n",
    "\n",
    "losses = []\n",
    "batches = len(train_loader)\n",
    "val_batches = len(val_loader)\n",
    "\n",
    "# loop for every epoch (training + evaluation)\n",
    "for epoch in range(epochs):\n",
    "    total_loss = 0\n",
    "\n",
    "    # progress bar (works in Jupyter notebook too!)\n",
    "    progress = tqdm(enumerate(train_loader), desc=\"Loss: \", total=batches)\n",
    "\n",
    "    # ----------------- TRAINING  -------------------- \n",
    "    # set model to training\n",
    "    model.train()\n",
    "    \n",
    "    for i, data in progress:\n",
    "        X, y = data[0].to(device), data[1].to(device)\n",
    "        # training step for single batch\n",
    "        model.zero_grad() # to make sure that all the grads are 0 \n",
    "        \"\"\"\n",
    "        model.zero_grad() and optimizer.zero_grad() are the same \n",
    "        IF all your model parameters are in that optimizer. \n",
    "        I found it is safer to call model.zero_grad() to make sure all grads are zero, \n",
    "        e.g. if you have two or more optimizers for one model.\n",
    "\n",
    "        \"\"\"\n",
    "        outputs = model(X) # forward\n",
    "        loss = loss_function(outputs, y) # get loss\n",
    "        loss.backward() # accumulates the gradient (by addition) for each parameter.\n",
    "        optimizer.step() # performs a parameter update based on the current gradient \n",
    "\n",
    "        # getting training quality data\n",
    "        current_loss = loss.item()\n",
    "        total_loss += current_loss\n",
    "\n",
    "        # updating progress bar\n",
    "        progress.set_description(\"Loss: {:.4f}\".format(total_loss/(i+1)))\n",
    "        \n",
    "    # releasing unceseccary memory in GPU\n",
    "    if torch.cuda.is_available():\n",
    "        torch.cuda.empty_cache()\n",
    "    \n",
    "    # ----------------- VALIDATION  ----------------- \n",
    "    val_losses = 0\n",
    "    precision, recall, f1, accuracy = [], [], [], []\n",
    "    \n",
    "    # set model to evaluating (testing)\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for i, data in enumerate(val_loader):\n",
    "            X, y = data[0].to(device), data[1].to(device)\n",
    "\n",
    "            outputs = model(X) # this get's the prediction from the network\n",
    "\n",
    "            val_losses += loss_function(outputs, y)\n",
    "\n",
    "            predicted_classes = torch.max(outputs, 1)[1] # get class from network's prediction\n",
    "            \n",
    "            # calculate P/R/F1/A metrics for batch\n",
    "            for acc, metric in zip((precision, recall, f1, accuracy), \n",
    "                                   (precision_score, recall_score, f1_score, accuracy_score)):\n",
    "                acc.append(\n",
    "                    calculate_metric(metric, y.cpu(), predicted_classes.cpu())\n",
    "                )\n",
    "          \n",
    "    print(f\"Epoch {epoch+1}/{epochs}, training loss: {total_loss/batches}, validation loss: {val_losses/val_batches}\")\n",
    "    print_scores(precision, recall, f1, accuracy, val_batches)\n",
    "    losses.append(total_loss/batches) # for plotting learning curve\n",
    "print(f\"Training time: {time.time()-start_ts}s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> ### 1.4 save model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'mnist_state.pt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> # 2. Data Generation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> ### 2.1 load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(60000, 28, 28)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train = x_train[:30001]\n",
    "y_train = y_train[:30001]\n",
    "x_test = x_test[:9000]\n",
    "y_test = y_test[:9000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x_train shape: (30001, 28, 28)\n",
      "Number of images in x_train 30001\n",
      "Number of images in x_test 9000\n"
     ]
    }
   ],
   "source": [
    "# Making sure that the values are float so that we can get decimal points after division\n",
    "x_train = x_train.astype('float32')\n",
    "x_test = x_test.astype('float32')\n",
    "\n",
    "# Normalizing the RGB codes by dividing it to the max RGB value.\n",
    "x_train /= 255\n",
    "x_test /= 255\n",
    "print('x_train shape:', x_train.shape)\n",
    "print('Number of images in x_train', x_train.shape[0])\n",
    "print('Number of images in x_test', x_test.shape[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> ### 2.2 Create tuple (index, label) for train and test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "instance_index_label = [(i, y_train[i]) for i in range(x_train.shape[0])]\n",
    "instance_index_label_test = [(i, y_test[i]) for i in range(x_test.shape[0])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# find the index if label is 1\n",
    "find_index = [instance_index_label[i][0] for i in range(len(instance_index_label)) if instance_index_label[i][1]==1]\n",
    "# find the index if label is 1\n",
    "find_index_test = [instance_index_label_test[i][0] for i in range(len(instance_index_label_test))\n",
    "                   if instance_index_label_test[i][1]==1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "index: 0\n",
      "label: 5\n"
     ]
    }
   ],
   "source": [
    "print('index:', instance_index_label[0][0]) #index\n",
    "print('label:', instance_index_label[0][1]) #label"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> ### 2.3 load pretrained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchvision.models.resnet import ResNet, BasicBlock\n",
    "class MnistResNet(ResNet):\n",
    "    def __init__(self):\n",
    "        super(MnistResNet, self).__init__(BasicBlock, [2, 2, 2, 2], num_classes=10)\n",
    "        self.conv1 = torch.nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return torch.softmax(super(MnistResNet, self).forward(x), dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): Conv2d(1, 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): BasicBlock(\n",
       "      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (5): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (6): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (7): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=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): BasicBlock(\n",
       "      (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (8): AdaptiveAvgPool2d(output_size=(1, 1))\n",
       ")"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = MnistResNet()\n",
    "model.load_state_dict(torch.load('mnist_state.pt'))\n",
    "body = nn.Sequential(*list(model.children()))\n",
    "# extract the last layer\n",
    "model = body[:9]\n",
    "# the model we will use\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> ### 2.4 get features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/linqisheng/anaconda3/lib/python3.7/site-packages/torchvision/datasets/mnist.py:53: UserWarning: train_data has been renamed data\n",
      "  warnings.warn(\"train_data has been renamed data\")\n"
     ]
    }
   ],
   "source": [
    "train_batch_size = 1\n",
    "val_batch_size = 1\n",
    "train_loader, val_loader = get_data_loaders(train_batch_size, val_batch_size)\n",
    "loss_function = nn.CrossEntropyLoss() # your loss function, cross entropy works well for multi-class problems\n",
    "\n",
    "# optimizer\n",
    "optimizer = optim.Adadelta(model.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "losses = []\n",
    "batches = len(train_loader)\n",
    "val_batches = len(val_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> #### 2.4.1 get features for train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c73c6c5955914d84ac3c5c2ddd9fc002",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Loss: ', max=60000, style=ProgressStyle(description_width='in…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# loop for every epoch (training + evaluation)\n",
    "meta_table = dict()\n",
    "feature_result = []\n",
    "\n",
    "# progress bar\n",
    "progress = tqdm(enumerate(train_loader), desc=\"Loss: \", total=batches)\n",
    "\n",
    "model.eval()\n",
    "\n",
    "for i, data in progress:\n",
    "    if i==30001:\n",
    "        break\n",
    "    X, y = data[0], data[1]\n",
    "    # training step for single batch\n",
    "    model.zero_grad()\n",
    "    outputs = model(X)\n",
    "    feature_result.append(outputs.reshape(-1).tolist())\n",
    "    meta_table[i] = outputs.reshape(-1).tolist()\n",
    "    \n",
    "feature_array = np.array(feature_result)\n",
    "np.save('feature_array_full',feature_array )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load\n",
    "feature_array = np.load('feature_array.npy', allow_pickle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> #### 2.4.2 get features for test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0dd0291eeaf341bbb29f46744214cc80",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Loss: ', max=60000, style=ProgressStyle(description_width='in…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# loop for every epoch (training + evaluation)\n",
    "meta_t_table = dict()\n",
    "feature_t_result = []\n",
    "\n",
    "# progress bar\n",
    "progress = tqdm(enumerate(val_loader), desc=\"Loss: \", total=batches)\n",
    "\n",
    "model.eval()\n",
    "\n",
    "for i, data in progress:\n",
    "    if i==9000:\n",
    "        break\n",
    "    X, y = data[0], data[1]\n",
    "    # training step for single batch\n",
    "    model.zero_grad()\n",
    "    outputs_t = model(X)\n",
    "    feature_t_result.append(outputs_t.reshape(-1).tolist())\n",
    "    meta_t_table[i] = outputs_t.reshape(-1).tolist()\n",
    "\n",
    "feature_test_array = np.array(feature_t_result)\n",
    "# save \n",
    "np.save('feature_test_array_full',feature_test_array )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load\n",
    "feature_test_array = np.load('feature_t_array.npy', allow_pickle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> ### 2.5 generate data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> #### 2.5.1 generate data for train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Dict, Tuple\n",
    "def data_generation(instance_index_label: List[Tuple]) -> List[Dict]:\n",
    "    \"\"\"\n",
    "    bags: {key1: [ind1, ind2, ind3],\n",
    "           key2: [ind1, ind2, ind3, ind4, ind5],\n",
    "           ... }\n",
    "    bag_lbls:\n",
    "        {key1: 0,\n",
    "         key2: 1,\n",
    "         ... }\n",
    "    \"\"\"\n",
    "    bag_size = np.random.randint(3,7,size=len(instance_index_label)//5)\n",
    "    data_cp = copy.copy(instance_index_label)\n",
    "    np.random.shuffle(data_cp)\n",
    "    bags = {}\n",
    "    bags_per_instance_labels = {}\n",
    "    bags_labels = {}\n",
    "    for bag_ind, size in enumerate(bag_size):\n",
    "        bags[bag_ind] = []\n",
    "        bags_per_instance_labels[bag_ind] = []\n",
    "        try:\n",
    "            for _ in range(size):\n",
    "                inst_ind, lbl = data_cp.pop()\n",
    "                bags[bag_ind].append(inst_ind)\n",
    "                # simplfy, just use a temporary variable instead of bags_per_instance_labels\n",
    "                bags_per_instance_labels[bag_ind].append(lbl)\n",
    "            bags_labels[bag_ind] = bag_label_from_instance_labels(bags_per_instance_labels[bag_ind])\n",
    "        except:\n",
    "            break\n",
    "    return bags, bags_labels\n",
    "\n",
    "def bag_label_from_instance_labels(instance_labels):\n",
    "    return int(any(((x==1) for x in instance_labels)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "bag_indices, bag_labels = data_generation(instance_index_label)\n",
    "bag_features = {kk: torch.Tensor(feature_array[inds]) for kk, inds in bag_indices.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save\n",
    "import pickle\n",
    "pickle.dump(bag_indices, open( \"bag_indices\", \"wb\" ) )\n",
    "pickle.dump(bag_labels, open( \"bag_labels\", \"wb\" ) )\n",
    "pickle.dump(bag_features, open( \"bag_features\", \"wb\" ) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "bag_indices = pickle.load( open( \"bag_indices\", \"rb\" ) )\n",
    "bag_labels = pickle.load( open( \"bag_labels\", \"rb\" ) )\n",
    "bag_features = pickle.load( open( \"bag_features\", \"rb\" ) )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> #### 2.5.2 generate data for test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "bag_t_indices, bag_t_labels = data_generation(instance_index_label_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "bag_t_features = {kk: torch.Tensor(feature_test_array[inds]) for kk, inds in bag_t_indices.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "pickle.dump(bag_t_indices, open( \"bag_t_indices\", \"wb\" ) )\n",
    "pickle.dump(bag_t_labels, open( \"bag_t_labels\", \"wb\" ) )\n",
    "pickle.dump(bag_t_features, open( \"bag_t_features\", \"wb\" ) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "bag_t_indices = pickle.load( open( \"bag_t_indices\", \"rb\" ) )\n",
    "bag_t_labels = pickle.load( open( \"bag_t_labels\", \"rb\" ) )\n",
    "bag_t_features = pickle.load( open( \"bag_t_features\", \"rb\" ) )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> # 3. Multiple Instance Learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> ### 3.1 Prepare data for model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset\n",
    "class Transform_data(Dataset):\n",
    "    \"\"\"\n",
    "    We want to 1. pad tensor 2. transform the data to the size that fits in the input size.\n",
    "    \n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, data, transform=None):\n",
    "        self.transform = transform\n",
    "        self.data = data\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        tensor = self.data[index][0]\n",
    "        if self.transform is not None:\n",
    "            tensor = self.transform(tensor)\n",
    "        return (tensor, self.data[index][1])\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = [(bag_features[i],bag_labels[i]) for i in range(len(bag_features))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[2.2300, 1.2076, 0.4186,  ..., 0.4665, 1.6642, 2.2301],\n",
       "        [0.8595, 1.5297, 0.9046,  ..., 0.5341, 0.9438, 0.8581],\n",
       "        [2.6144, 0.2394, 0.3311,  ..., 0.0778, 1.6823, 1.7735],\n",
       "        [0.1552, 1.2130, 1.0170,  ..., 2.8621, 0.8361, 0.2522],\n",
       "        [0.7179, 0.3636, 0.4274,  ..., 0.2156, 0.3240, 0.0769]])"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bag_features[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pad_tensor(data:list, max_number_instance) -> list:\n",
    "    \"\"\"\n",
    "    Since our bag has different sizes, we need to pad each tensor to have the same shape (max: 7).\n",
    "    We will look through each one instance and look at the shape of the tensor, and then we will pad 7-n \n",
    "    to the existing tensor where n is the number of instances in the bag.\n",
    "    The function will return a padded data set.\"\"\"\n",
    "    new_data = []\n",
    "    for bag_index in range(len(data)):\n",
    "        tensor_size = len(data[bag_index][0])\n",
    "        pad_size = max_number_instance - tensor_size\n",
    "        p2d = (0,0, 0, pad_size)\n",
    "        padded = nn.functional.pad(data[bag_index][0], p2d, 'constant', 0)\n",
    "        new_data.append((padded, data[bag_index][1]))\n",
    "    return new_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_number_instance = 7\n",
    "padded_train = pad_tensor(train_data, max_number_instance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data = [(bag_t_features[i],bag_t_labels[i]) for i in range(len(bag_t_features))]\n",
    "padded_test = pad_tensor(test_data, max_number_instance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data_loaders(train_data, test_data, train_batch_size, val_batch_size):\n",
    "    train_loader = DataLoader(train_data, batch_size=train_batch_size, shuffle=True)\n",
    "    val_loader = DataLoader(test_data, batch_size=val_batch_size, shuffle=False)\n",
    "    return train_loader, val_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader,valid_loader = get_data_loaders(padded_train, padded_test, 1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_batch_size = 1\n",
    "val_batch_size = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> ### 3.2 Define model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> #### 3.2.1 Aggregation Funtion model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# aggregation functions\n",
    "\n",
    "class SoftMaxMeanSimple(torch.nn.Module):\n",
    "    def __init__(self, n, n_inst, dim=0):\n",
    "        \"\"\"\n",
    "        if dim==1:\n",
    "            given a tensor `x` with dimensions [N * M],\n",
    "            where M -- dimensionality of the featur vector\n",
    "                       (number of features per instance)\n",
    "                  N -- number of instances\n",
    "            initialize with `AggModule(M)`\n",
    "            returns:\n",
    "            - weighted result: [M]\n",
    "            - gate: [N]\n",
    "        if dim==0:\n",
    "            ...\n",
    "        \"\"\"\n",
    "        super(SoftMaxMeanSimple, self).__init__()\n",
    "        self.dim = dim\n",
    "        self.gate = torch.nn.Softmax(dim=self.dim)      \n",
    "        self.mdl_instance_transform = nn.Sequential(\n",
    "                            nn.Linear(n, n_inst),\n",
    "                            nn.LeakyReLU(),\n",
    "                            nn.Linear(n_inst, n),\n",
    "                            nn.LeakyReLU(),\n",
    "                            )\n",
    "    def forward(self, x):\n",
    "        z = self.mdl_instance_transform(x)\n",
    "        if self.dim==0:\n",
    "            z = z.view((z.shape[0],1)).sum(1)\n",
    "        elif self.dim==1:\n",
    "            z = z.view((1, z.shape[1])).sum(0)\n",
    "        gate_ = self.gate(z)\n",
    "        res = torch.sum(x* gate_, self.dim)\n",
    "        return res, gate_\n",
    "\n",
    "    \n",
    "class AttentionSoftMax(torch.nn.Module):\n",
    "    def __init__(self, in_features = 3, out_features = None):\n",
    "        \"\"\"\n",
    "        given a tensor `x` with dimensions [N * M],\n",
    "        where M -- dimensionality of the featur vector\n",
    "                   (number of features per instance)\n",
    "              N -- number of instances\n",
    "        initialize with `AggModule(M)`\n",
    "        returns:\n",
    "        - weighted result: [M]\n",
    "        - gate: [N]\n",
    "        \"\"\"\n",
    "        super(AttentionSoftMax, self).__init__()\n",
    "        self.otherdim = ''\n",
    "        if out_features is None:\n",
    "            out_features = in_features\n",
    "        self.layer_linear_tr = nn.Linear(in_features, out_features)\n",
    "        self.activation = nn.LeakyReLU()\n",
    "        self.layer_linear_query = nn.Linear(out_features, 1)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        keys = self.layer_linear_tr(x)\n",
    "        keys = self.activation(keys)\n",
    "        attention_map_raw = self.layer_linear_query(keys)[...,0]\n",
    "        attention_map = nn.Softmax(dim=-1)(attention_map_raw)\n",
    "        result = torch.einsum(f'{self.otherdim}i,{self.otherdim}ij->{self.otherdim}j', attention_map, x)\n",
    "        return result, attention_map\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> #### 3.3.2 MIL_NN model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NoisyAnd(torch.nn.Module):\n",
    "    def __init__(self, a=10, dims=[1,2]):\n",
    "        super(NoisyAnd, self).__init__()\n",
    "#         self.output_dim = output_dim\n",
    "        self.a = a\n",
    "        self.b = torch.nn.Parameter(torch.tensor(0.01))\n",
    "        self.dims =dims\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "    def forward(self, x):\n",
    "#         h_relu = self.linear1(x).clamp(min=0)\n",
    "        mean = torch.mean(x, self.dims, True)\n",
    "        res = (self.sigmoid(self.a * (mean - self.b)) - self.sigmoid(-self.a * self.b)) / (\n",
    "              self.sigmoid(self.a * (1 - self.b)) - self.sigmoid(-self.a * self.b))\n",
    "        return res\n",
    "    \n",
    "\n",
    "\n",
    "class NN(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, n=512, n_mid = 1024,\n",
    "                 n_out=1, dropout=0.2,\n",
    "                 scoring = None,\n",
    "                ):\n",
    "        super(NN, self).__init__()\n",
    "        self.linear1 = torch.nn.Linear(n, n_mid)\n",
    "        self.non_linearity = torch.nn.LeakyReLU()\n",
    "        self.linear2 = torch.nn.Linear(n_mid, n_out)\n",
    "        self.dropout = torch.nn.Dropout(dropout)\n",
    "        if scoring:\n",
    "            self.scoring = scoring\n",
    "        else:\n",
    "            self.scoring = torch.nn.Softmax() if n_out>1 else torch.nn.Sigmoid()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        z = self.linear1(x)\n",
    "        z = self.non_linearity(z)\n",
    "        z = self.dropout(z)\n",
    "        z = self.linear2(z)\n",
    "        y_pred = self.scoring(z)\n",
    "        return y_pred\n",
    "    \n",
    "\n",
    "class LogisticRegression(torch.nn.Module):\n",
    "    def __init__(self, n=512, n_out=1):\n",
    "        super(LogisticRegression, self).__init__()\n",
    "        self.linear = torch.nn.Linear(n, n_out)\n",
    "        self.scoring = torch.nn.Softmax() if n_out>1 else torch.nn.Sigmoid()\n",
    "\n",
    "    def forward(self, x):\n",
    "        z = self.linear(x)\n",
    "        y_pred = self.scoring(z)\n",
    "        return y_pred\n",
    "\n",
    "    \n",
    "def regularization_loss(params,\n",
    "                        reg_factor = 0.005,\n",
    "                        reg_alpha = 0.5):\n",
    "    params = [pp for pp in params if len(pp.shape)>1]\n",
    "    l1_reg = nn.L1Loss()\n",
    "    l2_reg = nn.MSELoss()\n",
    "    loss_reg =0\n",
    "    for pp in params:\n",
    "        loss_reg+=reg_factor*((1-reg_alpha)*l1_reg(pp, target=torch.zeros_like(pp)) +\\\n",
    "                           reg_alpha*l2_reg(pp, target=torch.zeros_like(pp)))\n",
    "    return loss_reg\n",
    "\n",
    "class MIL_NN(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, n=512,  \n",
    "                 n_mid=1024, \n",
    "                 n_classes=1, \n",
    "                 dropout=0.1,\n",
    "                 agg = None,\n",
    "                 scoring=None,\n",
    "                ):\n",
    "        super(MIL_NN, self).__init__()\n",
    "        self.agg = agg if agg is not None else AttentionSoftMax(n)\n",
    "        \n",
    "        if n_mid == 0:\n",
    "            self.bag_model = LogisticRegression(n, n_classes)\n",
    "        else:\n",
    "            self.bag_model = NN(n, n_mid, n_classes, dropout=dropout, scoring=scoring)\n",
    "        \n",
    "    def forward(self, bag_features, bag_lbls=None):\n",
    "        \"\"\"\n",
    "        bag_feature is an aggregated vector of 512 features\n",
    "        bag_att is a gate vector of n_inst instances\n",
    "        bag_lbl is a vector a labels\n",
    "        figure out batches\n",
    "        \"\"\"\n",
    "        bag_feature, bag_att, bag_keys = list(zip(*[list(self.agg(ff.float())) + [idx]\n",
    "                                                    for idx, ff in (bag_features.items())]))\n",
    "        bag_att = dict(zip(bag_keys, [a.detach().cpu() for a  in bag_att]))\n",
    "        bag_feature_stacked = torch.stack(bag_feature)\n",
    "        y_pred = self.bag_model(bag_feature_stacked)\n",
    "        return y_pred, bag_att, bag_keys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> #### 3.3.3 helper function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_metric(metric_fn, true_y, pred_y):\n",
    "    # multi class problems need to have averaging method\n",
    "    if \"average\" in inspect.getfullargspec(metric_fn).args:\n",
    "        return metric_fn(true_y, pred_y, average=\"macro\")\n",
    "    else:\n",
    "        return metric_fn(true_y, pred_y)\n",
    "    \n",
    "def print_scores(p, r, f1, a, batch_size):\n",
    "    # just an utility printing function\n",
    "    for name, scores in zip((\"precision\", \"recall\", \"F1\", \"accuracy\"), (p, r, f1, a)):\n",
    "        print(f\"\\t{name.rjust(14, ' ')}: {sum(scores)/batch_size:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> # 4. Train & Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1b035bafd20f46bb911b51b90256767e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Loss: ', max=6000, style=ProgressStyle(description_width='ini…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/linqisheng/anaconda3/lib/python3.7/site-packages/torch/nn/modules/loss.py:498: UserWarning: Using a target size (torch.Size([1])) that is different to the input size (torch.Size([1, 1])) is deprecated. Please ensure they have the same size.\n",
      "  return F.binary_cross_entropy(input, target, weight=self.weight, reduction=self.reduction)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/linqisheng/anaconda3/lib/python3.7/site-packages/sklearn/metrics/classification.py:1437: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 in labels with no predicted samples.\n",
      "  'precision', 'predicted', average, warn_for)\n",
      "/home/linqisheng/anaconda3/lib/python3.7/site-packages/sklearn/metrics/classification.py:1439: UndefinedMetricWarning: Recall is ill-defined and being set to 0.0 in labels with no true samples.\n",
      "  'recall', 'true', average, warn_for)\n",
      "/home/linqisheng/anaconda3/lib/python3.7/site-packages/sklearn/metrics/classification.py:1437: UndefinedMetricWarning: F-score is ill-defined and being set to 0.0 in labels with no predicted samples.\n",
      "  'precision', 'predicted', average, warn_for)\n",
      "/home/linqisheng/anaconda3/lib/python3.7/site-packages/sklearn/metrics/classification.py:1439: UndefinedMetricWarning: F-score is ill-defined and being set to 0.0 in labels with no true samples.\n",
      "  'recall', 'true', average, warn_for)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10, training loss: 0.3253995268222628, validation loss: 0.17935876548290253\n",
      "\t     precision: 0.9411\n",
      "\t        recall: 0.9411\n",
      "\t            F1: 0.9411\n",
      "\t      accuracy: 0.9411\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "532909da90da493c92b7da30548c2ad3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Loss: ', max=6000, style=ProgressStyle(description_width='ini…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 2/10, training loss: 0.17075676332922501, validation loss: 0.1573299914598465\n",
      "\t     precision: 0.9361\n",
      "\t        recall: 0.9361\n",
      "\t            F1: 0.9361\n",
      "\t      accuracy: 0.9361\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "43ef20d1b85f40cd9ce454c092ced61e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Loss: ', max=6000, style=ProgressStyle(description_width='ini…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 3/10, training loss: 0.1241153416950585, validation loss: 0.14403362572193146\n",
      "\t     precision: 0.9461\n",
      "\t        recall: 0.9461\n",
      "\t            F1: 0.9461\n",
      "\t      accuracy: 0.9461\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d61c59ea4bbc4507afc1092d19cec538",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Loss: ', max=6000, style=ProgressStyle(description_width='ini…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 4/10, training loss: 0.10115043728419369, validation loss: 0.11335857212543488\n",
      "\t     precision: 0.9628\n",
      "\t        recall: 0.9628\n",
      "\t            F1: 0.9628\n",
      "\t      accuracy: 0.9628\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4e28f49ff04444839423c5d3b0fc3798",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Loss: ', max=6000, style=ProgressStyle(description_width='ini…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 5/10, training loss: 0.07484548372746803, validation loss: 0.07197500765323639\n",
      "\t     precision: 0.9778\n",
      "\t        recall: 0.9778\n",
      "\t            F1: 0.9778\n",
      "\t      accuracy: 0.9778\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "76e917943e824e6386b41b2bb2ee0713",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Loss: ', max=6000, style=ProgressStyle(description_width='ini…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 6/10, training loss: 0.060945956385883286, validation loss: 0.07357295602560043\n",
      "\t     precision: 0.9739\n",
      "\t        recall: 0.9739\n",
      "\t            F1: 0.9739\n",
      "\t      accuracy: 0.9739\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ecfad9da632e40179123ba79a1c17ada",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Loss: ', max=6000, style=ProgressStyle(description_width='ini…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 7/10, training loss: 0.05318203048250962, validation loss: 0.074022576212883\n",
      "\t     precision: 0.9756\n",
      "\t        recall: 0.9756\n",
      "\t            F1: 0.9756\n",
      "\t      accuracy: 0.9756\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a07cd1f673b24579a2369deef71f89d5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Loss: ', max=6000, style=ProgressStyle(description_width='ini…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 8/10, training loss: 0.04461347977752607, validation loss: 0.0600493848323822\n",
      "\t     precision: 0.9783\n",
      "\t        recall: 0.9783\n",
      "\t            F1: 0.9783\n",
      "\t      accuracy: 0.9783\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "41b52fa7f08d4b71ab5286a75eb93c9d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Loss: ', max=6000, style=ProgressStyle(description_width='ini…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 9/10, training loss: 0.03741783676902158, validation loss: 0.0561109222471714\n",
      "\t     precision: 0.9856\n",
      "\t        recall: 0.9856\n",
      "\t            F1: 0.9856\n",
      "\t      accuracy: 0.9856\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0d733ee1964c4ef29af11927bbb9ac5a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Loss: ', max=6000, style=ProgressStyle(description_width='ini…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 10/10, training loss: 0.029147379207246065, validation loss: 0.05226973071694374\n",
      "\t     precision: 0.9850\n",
      "\t        recall: 0.9850\n",
      "\t            F1: 0.9850\n",
      "\t      accuracy: 0.9850\n",
      "Training time: 690.9660251140594s\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "start_ts = time.time()\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "lr0 = 1e-4\n",
    "\n",
    "# model:\n",
    "model = MIL_NN().to(device)\n",
    "\n",
    "# params you need to specify:\n",
    "epochs = 10\n",
    "train_loader, val_loader = get_data_loaders(padded_train, padded_test, 1, 1)\n",
    "loss_function = torch.nn.BCELoss(reduction='mean') # your loss function, cross entropy works well for multi-class problems\n",
    "\n",
    "\n",
    "#optimizer = optim.Adadelta(model.parameters())\n",
    "optimizer = optim.SGD(model.parameters(), lr=lr0, momentum=0.9)\n",
    "\n",
    "losses = []\n",
    "batches = len(train_loader)\n",
    "val_batches = len(val_loader)\n",
    "\n",
    "# loop for every epoch (training + evaluation)\n",
    "for epoch in range(epochs):\n",
    "    total_loss = 0\n",
    "\n",
    "    # progress bar (works in Jupyter notebook too!)\n",
    "    progress = tqdm(enumerate(train_loader), desc=\"Loss: \", total=batches)\n",
    "\n",
    "    # ----------------- TRAINING  -------------------- \n",
    "    # set model to training\n",
    "    model.train()\n",
    "    for i, data in progress:\n",
    "        X, y = data[0].to(device), data[1].to(device)\n",
    "        X = X.reshape([1,7*512])\n",
    "        y = y.type(torch.cuda.FloatTensor)\n",
    "        # training step for single batch\n",
    "        model.zero_grad() # to make sure that all the grads are 0 \n",
    "        \"\"\"\n",
    "        model.zero_grad() and optimizer.zero_grad() are the same \n",
    "        IF all your model parameters are in that optimizer. \n",
    "        I found it is safer to call model.zero_grad() to make sure all grads are zero, \n",
    "        e.g. if you have two or more optimizers for one model.\n",
    "\n",
    "        \"\"\"\n",
    "        outputs = model(X) # forward\n",
    "        loss = loss_function(outputs, y) # get loss\n",
    "        loss.backward() # accumulates the gradient (by addition) for each parameter.\n",
    "        optimizer.step() # performs a parameter update based on the current gradient \n",
    "\n",
    "        # getting training quality data\n",
    "        current_loss = loss.item()\n",
    "        total_loss += current_loss\n",
    "\n",
    "        # updating progress bar\n",
    "        progress.set_description(\"Loss: {:.4f}\".format(total_loss/(i+1)))\n",
    "        \n",
    "    # releasing unceseccary memory in GPU\n",
    "    if torch.cuda.is_available():\n",
    "        torch.cuda.empty_cache()\n",
    "    \n",
    "    # ----------------- VALIDATION  ----------------- \n",
    "    val_losses = 0\n",
    "    precision, recall, f1, accuracy = [], [], [], []\n",
    "    \n",
    "    # set model to evaluating (testing)\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for i, data in enumerate(val_loader):\n",
    "            X, y = data[0].to(device), data[1].to(device)\n",
    "            X = X.reshape([1,7*512])\n",
    "            y = y.type(torch.cuda.FloatTensor)\n",
    "            outputs = model(X) # this get's the prediction from the network\n",
    "            prediced_classes =outputs.detach().round()\n",
    "            #y_pred.extend(prediced_classes.tolist())\n",
    "            val_losses += loss_function(outputs, y)\n",
    "            \n",
    "            # calculate P/R/F1/A metrics for batch\n",
    "            for acc, metric in zip((precision, recall, f1, accuracy), \n",
    "                                   (precision_score, recall_score, f1_score, accuracy_score)):\n",
    "                acc.append(\n",
    "                    calculate_metric(metric, y.cpu(), prediced_classes.cpu())\n",
    "                )\n",
    "          \n",
    "    print(f\"Epoch {epoch+1}/{epochs}, training loss: {total_loss/batches}, validation loss: {val_losses/val_batches}\")\n",
    "    print_scores(precision, recall, f1, accuracy, val_batches)\n",
    "    losses.append(total_loss/batches) # for plotting learning curve\n",
    "print(f\"Training time: {time.time()-start_ts}s\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}