MAIN_ABMIL_with_within_bag_sampling_IMAGENETTE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code is based on implementation (https://github.com/AMLab-Amsterdam/AttentionDeepMIL) of the paper [1]. The changes are made according to the paper [2]. <br>\n",
    "[1] Ilse, Maximilian, Jakub Tomczak, and Max Welling \"Attention-based deep multiple instance learning.\" International conference on machine learning. PMLR, 2018 <br>\n",
    "[2] \"The Effect of Within-Bag Sampling on End-to-End Multiple Instance Learning\", 2021"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# External dependencies\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.utils.data as data_utils\n",
    "from torchvision import datasets, transforms\n",
    "import argparse\n",
    "import torch.optim as optim\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable\n",
    "import os\n",
    "import shutil\n",
    "from pathlib import Path\n",
    "from PIL import Image\n",
    "# Internal modules\n",
    "import attention_model\n",
    "from dataloaders import ImagenetteBags\n",
    "from attention_model import Attention_Imagenette_bags_3GPUs, Attention_Imagenette_bags_1GPU\n",
    "from evaluation import compute_metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def list_files_in_folder(image_folder):\n",
    "    \"\"\"Lists file names in a given directory\"\"\"\n",
    "    list_of_files = []\n",
    "    for file in os.listdir(image_folder):\n",
    "        if os.path.isfile(os.path.join(image_folder, file)):\n",
    "            list_of_files.append(file)\n",
    "    return list_of_files\n",
    "\n",
    "def create_save_dir(direct, name_subdirectory):\n",
    "    if not os.path.exists(os.path.join(direct, name_subdirectory)):\n",
    "        print('make dir')\n",
    "        os.mkdir(os.path.join(direct, name_subdirectory))\n",
    "    return os.path.join(direct, name_subdirectory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training settings\n",
    "parser = argparse.ArgumentParser(description='PyTorch Imegenette bags Example')\n",
    "parser.add_argument('--epochs', type=int, default=50, metavar='N',\n",
    "                    help='number of epochs to train (default: 20)')\n",
    "parser.add_argument('--lr', type=float, default=0.0001, metavar='LR',\n",
    "                    help='learning rate (default: 0.0005)')\n",
    "parser.add_argument('--reg', type=float, default=10e-5, metavar='R',\n",
    "                    help='weight decay')\n",
    "parser.add_argument('--num_bags_test', type=int, default=10, metavar='NTest',\n",
    "                    help='number of bags in test set')\n",
    "parser.add_argument('--seed', type=int, default=1, metavar='S',\n",
    "                    help='random seed (default: 1)')\n",
    "\n",
    "# args = parser.parse_args()\n",
    "args, unknown = parser.parse_known_args()\n",
    "\n",
    "data_path = './fold1/Imagenette_0010_0050_0030'\n",
    "path = Path(data_path)\n",
    "mean_data = [0.485, 0.456, 0.406]\n",
    "std_data = [0.229, 0.224, 0.225]\n",
    "# Parameters derived from dataset name (as in \"Create_Imagenette_bags_dataset.ipynb\"), otherwise enter\n",
    "num_bags_train = int(path.parts[-1][11:15])\n",
    "mean_bag_length = int(path.parts[-1][16:-5])\n",
    "print(num_bags_train,mean_bag_length)\n",
    "# Three or one GPU: three_gpus is True if two GPUs are used, False if one\n",
    "three_gpus = False\n",
    "args.sampling_percent = 30\n",
    "image_sizeImagenette = (3,112,112)\n",
    "\n",
    "sampling_size_in_instances = np.ceil((args.sampling_percent*mean_bag_length)/100)\n",
    "N_repeat = 10 # Approximate number of times one image is sampled to bags when sampling percent is lower than 100\n",
    "if args.sampling_percent<100:\n",
    "    test_epochs = np.ceil((N_repeat*mean_bag_length)/sampling_size_in_instances).astype(np.int)\n",
    "    val_times = 5 # The number of times validation is repeated (with resampling of images from the bags) \n",
    "else:\n",
    "    test_epochs = 1; val_times = 1\n",
    "window_length = 15 # the moving average window length to find validation model with lowest error\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args.cuda = torch.cuda.is_available()\n",
    "torch.manual_seed(args.seed)\n",
    "if args.cuda:\n",
    "    torch.cuda.manual_seed(args.seed)\n",
    "    print('\\nGPU is ON!')\n",
    "print('Init Model')\n",
    "if args.cuda:\n",
    "    if three_gpus:\n",
    "        model = Attention_Imagenette_bags_3GPUs() \n",
    "    else:\n",
    "        model = Attention_Imagenette_bags_1GPU()\n",
    "        model.to('cuda:0')\n",
    "save_weights_dir = create_save_dir(data_path, 'test_weights_epochs_'+str(test_epochs)+'samplingPerc_'+\n",
    "                                   str(args.sampling_percent)+'_train_epochs_'+str(args.epochs)+'_lr'+str(args.lr)) \n",
    "if os.listdir(save_weights_dir):\n",
    "    print('Directory is not empty!')\n",
    "optimizer = optim.Adam(model.parameters(), lr=args.lr, betas=(0.9, 0.999), weight_decay=args.reg)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = data_utils.DataLoader(ImagenetteBags(train=True,\n",
    "                                                    valid=False,\n",
    "                                                    test=False,\n",
    "                                                    image_size=image_sizeImagenette,\n",
    "                                     transform=transforms.Compose([transforms.ToTensor(),\n",
    "                                                                   transforms.Normalize(mean=mean_data, \n",
    "                                                                                        std=std_data)\n",
    "                                     ]), sampling_size=sampling_size_in_instances, data_path=data_path), \n",
    "                                     batch_size=1,\n",
    "                                     shuffle=True)\n",
    "\n",
    "valid_loader = data_utils.DataLoader(ImagenetteBags(train=False,\n",
    "                                                    valid=True,\n",
    "                                                    test=False,\n",
    "                                                    image_size=image_sizeImagenette,\n",
    "                                    transform=transforms.Compose([transforms.ToTensor(),\n",
    "                                                                  transforms.Normalize(mean=mean_data, \n",
    "                                                                                       std=std_data)\n",
    "                                    ]), sampling_size=sampling_size_in_instances, data_path=data_path),\n",
    "                                    batch_size=1,\n",
    "                                    shuffle=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = Path(data_path)\n",
    "all_train_loss = []; all_train_error=[]; all_valid_loss=[]; all_valid_error=[]\n",
    "count_stable_epochs=[]; \n",
    "if three_gpus:\n",
    "    min_valid_loss = torch.tensor([float('inf')]).to('cuda:1') \n",
    "else:\n",
    "    min_valid_loss = torch.tensor([float('inf')]).to('cuda:0') \n",
    "min_valid_error = np.inf; window_with_best_avg=np.inf\n",
    "num_epochs=args.epochs; lr=args.lr; sampling_percent=args.sampling_percent\n",
    "if not os.listdir(save_weights_dir): \n",
    "    for epoch in range(1, args.epochs + 1):\n",
    "        print(\"Training started\")\n",
    "        model.train()\n",
    "        train_loss = 0.; train_error = 0.\n",
    "        for batch_idx, (data, label, sample_indices, index, bages_names) in enumerate(train_loader):\n",
    "\n",
    "            bag_label = label[0]\n",
    "            if args.cuda:\n",
    "                if three_gpus:\n",
    "                    data, bag_label = data.to('cuda:2'), bag_label.to('cuda:1') \n",
    "                else:\n",
    "                    data, bag_label = data.to('cuda:0'), bag_label.to('cuda:0') \n",
    "            data, bag_label = Variable(data), Variable(bag_label)\n",
    "\n",
    "            # reset gradients\n",
    "            optimizer.zero_grad()\n",
    "            # calculate loss and metrics\n",
    "            loss, _ = model.calculate_objective(data, bag_label)\n",
    "            train_loss += loss.data[0]\n",
    "            error, _ = model.calculate_classification_error(data, bag_label)\n",
    "            train_error += error\n",
    "\n",
    "            # backward pass\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "        # calculate loss and error for epoch\n",
    "        train_loss /= len(train_loader)\n",
    "        train_error /= len(train_loader)\n",
    "\n",
    "        # Validation loss and error\n",
    "        valid_loss = 0.; valid_error = 0.\n",
    "        model.eval()\n",
    "        for epp in range(0, val_times):\n",
    "            for batch_idx, (data, label, sample_indices, index, bages_names) in enumerate(valid_loader):\n",
    "                bag_label = label[0]\n",
    "                if args.cuda:\n",
    "                    if three_gpus:\n",
    "                        data, bag_label = data.to('cuda:2'), bag_label.to('cuda:1') \n",
    "                    else:\n",
    "                        data, bag_label = data.to('cuda:0'), bag_label.to('cuda:0') \n",
    "                data, bag_label = Variable(data), Variable(bag_label)  \n",
    "                loss, _ = model.calculate_objective(data, bag_label)\n",
    "                valid_loss += loss.data[0]\n",
    "                error, _ = model.calculate_classification_error(data, bag_label)\n",
    "                valid_error += error\n",
    "        valid_loss /= len(valid_loader)*val_times\n",
    "        valid_error /= len(valid_loader)*val_times\n",
    "\n",
    "        # Save all losses for all epochs both valid and train\n",
    "        all_train_loss.append(train_loss.cpu().data.numpy())\n",
    "        all_train_error.append(train_error)\n",
    "        all_valid_loss.append(valid_loss.cpu().data.numpy())\n",
    "        all_valid_error.append(valid_error)\n",
    "\n",
    "        # Best validation error within the best average window of error\n",
    "        name_epoch_model = os.path.join(str(path), 'saved_model_IMAGENETTE_ResNet_'+str(path.parts[-1])+'_currentEpoch'+\\\n",
    "                                        str(epoch)+'_overallEpochs'+str(num_epochs)+'_lr'+str(lr)+'_sampPerc'+\\\n",
    "                                        str(sampling_percent)+'.pt')\n",
    "        torch.save(model.state_dict(), name_epoch_model)\n",
    "        if len(all_valid_error)>window_length-1:\n",
    "            avg_valid_error_window = np.mean(np.asarray(all_valid_error)[-window_length:])\n",
    "\n",
    "            if window_with_best_avg > avg_valid_error_window:\n",
    "                idx_model_in_window_with_min_loss = np.argmin(np.asarray(all_valid_error)[-window_length:])\n",
    "                window_with_best_avg = avg_valid_error_window\n",
    "                name_best_inMovAv_val_error_model = os.path.join(str(path), 'saved_model_IMAGENETTE_ResNet_'+str(path.parts[-1])+\n",
    "                                                                '_currentEpoch'+\n",
    "                                                 str(epoch-window_length+idx_model_in_window_with_min_loss+1)+\n",
    "                                                 '_overallEpochs'+str(num_epochs)+'_lr'+str(lr)+\n",
    "                                                 '_sampPerc'+str(sampling_percent)+'.pt')\n",
    "                print(\"Current best moving average model, epoch: \", str(epoch-window_length+idx_model_in_window_with_min_loss+1))\n",
    "                epoch_best_movingAvg = epoch-window_length+idx_model_in_window_with_min_loss+1\n",
    "        #to delete models that are outside current window except the best_window_model_name\n",
    "        for ep in range(0, epoch-window_length):\n",
    "            model_nametmp = os.path.join(str(path), 'saved_model_IMAGENETTE_ResNet_'+str(path.parts[-1])+'_currentEpoch'+\n",
    "                                         str(ep+1)+'_overallEpochs'+str(num_epochs)+'_lr'+str(lr)+'_sampPerc'+\n",
    "                                         str(sampling_percent)+'.pt')\n",
    "            if model_nametmp != name_best_inMovAv_val_error_model and os.path.isfile(model_nametmp):\n",
    "                os.remove(model_nametmp)\n",
    "\n",
    "        print('Epoch: {}, Loss: {:.4f}, Train error: {:.4f}'.format(epoch, train_loss.cpu().numpy()[0], train_error))\n",
    "        print('Epoch: {}, Loss: {:.4f}, Valid error: {:.4f}'.format(epoch, valid_loss.cpu().numpy()[0], valid_error))\n",
    "\n",
    "\n",
    "    new_name_best_inMovAv_val_error_model = os.path.join(save_weights_dir, 'saved_model_best_movingAvg_IMAGENETTE_ResNet_'+\\\n",
    "                                                         str(path.parts[-1])+'epochSaved_'+str(epoch_best_movingAvg)+\n",
    "                                                         '_overallEpochs'+str(num_epochs)+'_lr'+str(lr)+\n",
    "                                                         '_sampPerc'+str(sampling_percent)+'.pt')\n",
    "    shutil.move(name_best_inMovAv_val_error_model, new_name_best_inMovAv_val_error_model)\n",
    "\n",
    "\n",
    "    for ep in range(1,num_epochs+1):\n",
    "        model_nametmp = os.path.join(str(path), 'saved_model_IMAGENETTE_ResNet_'+str(path.parts[-1])+'_currentEpoch'+str(ep)+\n",
    "                                     '_overallEpochs'+str(num_epochs)+'_lr'+str(lr)+'_sampPerc'+str(sampling_percent)+'.pt')\n",
    "        if os.path.isfile(model_nametmp):\n",
    "            os.remove(model_nametmp)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# With replacement\n",
    "def test(save_weights_dir, subfolder, test_loader, model):\n",
    "    model.eval()\n",
    "    test_loss = 0.; test_error = 0.\n",
    "    for epoch in range(0, test_epochs): \n",
    "        print('test_epoch', epoch, end='\\r')\n",
    "        for batch_idx, (data, label, sample_indices, index, bages_names) in enumerate(test_loader):\n",
    "            bag_label = label[0]\n",
    "            instance_labels = label[1]\n",
    "            if args.cuda:\n",
    "                if three_gpus:\n",
    "                    data, bag_label = data.to('cuda:2'), bag_label.to('cuda:1') \n",
    "                else:\n",
    "                    data, bag_label = data.to('cuda:0'), bag_label.to('cuda:0') \n",
    "            data, bag_label = Variable(data), Variable(bag_label)\n",
    "            loss, attention_weights = model.calculate_objective(data, bag_label)\n",
    "            test_loss += loss.data[0]\n",
    "            error, predicted_label = model.calculate_classification_error(data, bag_label)\n",
    "            test_error += error\n",
    "            \n",
    "            create_save_dir(save_weights_dir, subfolder)\n",
    "            save_weights_bag = create_save_dir(os.path.join(save_weights_dir, subfolder), \\\n",
    "                                               str(bages_names[0].cpu().numpy()).zfill(4))\n",
    "            \n",
    "            for i in range(data.cpu().shape[1]):\n",
    "                label_name = int(label[1][0][i].cpu().numpy())\n",
    "                # Save attention weights\n",
    "                np.save(os.path.join(save_weights_bag, str(all_names_test[index][sample_indices[0][i]])+'_'+\\\n",
    "                                     str(epoch).zfill(2)+'_'+str(label_name)+'_'+\\\n",
    "                                     str(bages_names[0].cpu().numpy()).zfill(4)+'_'+\\\n",
    "                                     str(predicted_label.cpu().numpy())+\".npy\"), \\\n",
    "                        attention_weights.cpu().data.numpy()[0][i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "test_loader = data_utils.DataLoader(ImagenetteBags(train=False,\n",
    "                                                   valid=False,\n",
    "                                                   test=True,\n",
    "                                                   image_size=image_sizeImagenette,\n",
    "                                    transform=transforms.Compose([transforms.ToTensor(),\n",
    "                                                                  transforms.Normalize(mean=mean_data, \n",
    "                                                                                       std=std_data)\n",
    "                                    ]), sampling_size=sampling_size_in_instances, data_path=data_path),\n",
    "                                    batch_size=1,\n",
    "                                    shuffle=False)\n",
    "all_names_test = np.load(os.path.join(data_path,'test_imgs_lists.npy'), allow_pickle=True)\n",
    "    \n",
    "print('Start Testing')\n",
    "if args.cuda:\n",
    "    if three_gpus:\n",
    "        model = Attention_Imagenette_bags_3GPUs() \n",
    "    else:\n",
    "        model = Attention_Imagenette_bags_1GPU()\n",
    "        model.to('cuda:0')\n",
    "model.load_state_dict(torch.load(new_name_best_inMovAv_val_error_model))\n",
    "subfolder='best_inMovAv_val_loss_model_weights'\n",
    "if not os.listdir(save_weights_dir):\n",
    "    test(save_weights_dir, subfolder, test_loader, model)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_test_bags = args.num_bags_test\n",
    "compute_metrics('best_inMovAv_val_loss_model_weights', num_test_bags, save_weights_dir, data_path, test_epochs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}