ColorMNIST_OOD.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f632e357-f836-4d2e-af43-4a0519094ac0",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Data loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d316e7eb-971b-4364-b293-301ac6187200",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import torch.utils.data as data_utils\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "from dataloaders.dataloader import ColoredMnistBags3_binary, ColoredFashionMnistBags3_binary\n",
    "\n",
    "def mi_collate_img(batch):\n",
    "    # collate_fn for pytorch DataLoader\n",
    "    bag = [item[0] for item in batch]\n",
    "    bag = torch.tensor(np.concatenate(bag, axis = 0))\n",
    "    \n",
    "    bag_idx = [item[2] for item in batch]\n",
    "    bag_idx = torch.tensor(np.concatenate(bag_idx, axis = 0))\n",
    "    \n",
    "    label = [item[1] for item in batch]\n",
    "\n",
    "    instance_label = [item[1] for item in label]\n",
    "    instance_label = torch.tensor(np.concatenate(instance_label, axis = 0))\n",
    "    bag_label = [item[0] for item in label]\n",
    "    bag_label = torch.tensor(bag_label)\n",
    "    return bag, bag_idx, bag_label, instance_label"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9a47c8a-5e32-4fe3-a99a-61c6091d2c4d",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5426354f-6656-4e54-82f1-9c876ea95ab2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from util import *\n",
    "from models.networks_colormnist import cmil_mnist\n",
    "\n",
    "from sklearn.metrics import accuracy_score\n",
    "from sklearn.metrics import roc_curve\n",
    "def training_procedure(FLAGS, input_dim, dataloader):\n",
    "  device = torch.device('cuda') \n",
    "  model = cmil_mnist(FLAGS).to(device)\n",
    "  # model.apply(weights_init)\n",
    "  model.train()\n",
    "  auto_encoder_optimizer = torch.optim.AdamW(model.parameters(), \n",
    "                                  lr=FLAGS.initial_learning_rate, weight_decay=FLAGS.weight_decay)\n",
    "  best_loss = 1000000.\n",
    "\n",
    "  for epoch in range(0, FLAGS.end_epoch):\n",
    "      elbo_epoch = 0\n",
    "      recon_epoch = 0\n",
    "      y_epoch = 0\n",
    "      KL_ins_epoch = 0\n",
    "      for (i, batch) in enumerate(dataloader):\n",
    "          bag, bag_idx, bag_label, _ = batch # do not use instance label while training\n",
    "          auto_encoder_optimizer.zero_grad()            \n",
    "          elbo, class_y_loss, reconstruction_proba, KL_instance = \\\n",
    "              model.loss_function(bag.float().to(device), bag_idx.to(device), bag_label.to(device), epoch)\n",
    "          elbo.backward()\n",
    "          auto_encoder_optimizer.step()  \n",
    "            \n",
    "          elbo_epoch  += elbo\n",
    "          recon_epoch += reconstruction_proba\n",
    "          y_epoch += class_y_loss\n",
    "          KL_ins_epoch += KL_instance\n",
    "      elbo_epoch = elbo_epoch / (dataloader.__len__()/batch_size)\n",
    "      recon_epoch = recon_epoch / (dataloader.__len__()/batch_size)\n",
    "      y_epoch = y_epoch / (dataloader.__len__()/batch_size)\n",
    "      KL_ins_epoch = KL_ins_epoch / (dataloader.__len__()/batch_size)\n",
    "      \n",
    "      if ((epoch + 1) % 20 ==0):\n",
    "          print('Epoch #' + str(epoch+1) + '..............................................')\n",
    "          print(\"Train elbo  {:.5f}, recon_loss {:.5f}, y_loss {:.5f}\" \\\n",
    "                .format (elbo_epoch, recon_epoch,  y_epoch))\n",
    "          print(\"KL zx  {:.6f}\".format (KL_ins_epoch))\n",
    "      if elbo_epoch < best_loss:\n",
    "            best_loss = elbo_epoch\n",
    "            torch.save(model.state_dict(), '/home/temp_cmnist.pt')\n",
    "  \n",
    "  model.load_state_dict(torch.load('/home/temp_cmnist.pt'))\n",
    "  return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c6e0a5f-6454-4df5-b11b-5ca816e82920",
   "metadata": {},
   "outputs": [],
   "source": [
    "def weight_reset(m):\n",
    "   if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear) or isinstance(m, nn.ConvTranspose2d):\n",
    "     m.reset_parameters()\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    import argparse\n",
    "    import torchvision\n",
    "    import numpy as np\n",
    "    import pandas as pd\n",
    "    from sklearn.model_selection  import ParameterGrid\n",
    "    import matplotlib.pyplot as plt\n",
    "    param_grid = {'instance_dim': [24], 'aux_loss_multiplier_y': [1000]}\n",
    "    grid = ParameterGrid(param_grid)    \n",
    "    for params in grid:\n",
    "        parser = argparse.ArgumentParser()\n",
    "        parser.add_argument('--cuda', type=bool, default=True, help=\"run the following code on a GPU\")\n",
    "        parser.add_argument('--num_classes', type=int, default=2, help=\"number of classes on which the data set trained\")\n",
    "        parser.add_argument('--initial_learning_rate', type=float, default=1e-3, help=\"starting learning rate\")\n",
    "        parser.add_argument(\"--weight-decay\", default=1e-4, type=float)\n",
    "        parser.add_argument('--instance_dim', type=int, default=params['instance_dim'], help=\"dimension of instance factor latent space\")\n",
    "        parser.add_argument('--reconstruction_coef', type=float, default=1., help=\"coefficient for reconstruction term\")\n",
    "        parser.add_argument('--kl_divergence_coef', type=float, default=1, help=\"coefficient for instance KL-Divergence loss term\")\n",
    "        parser.add_argument('--aux_loss_multiplier_y', type=float, default=params['aux_loss_multiplier_y'])\n",
    "        parser.add_argument('--start_epoch', type=int, default=0, help=\"flag to set the starting epoch for training\")\n",
    "        parser.add_argument('--end_epoch', type=int, default=200, help=\"flag to indicate the final epoch of training\")\n",
    "        parser.add_argument('-w', '--warmup', type=int, default=0, metavar='N', help='number of epochs for warm-up. Set to 0 to turn warmup off.')\n",
    "        parser.add_argument('--in_channels', type=int, default=3, metavar='N', help='number of input channels, used for colored mnists')\n",
    "\n",
    "        FLAGS = parser.parse_args(args=[])\n",
    "\n",
    "        batch_size = 1024\n",
    "        bag_size = 10\n",
    "        \n",
    "#         task = 'fashion'\n",
    "        task = 'mnist'\n",
    "        print('Current Task is', task)\n",
    "        if task == 'mnist':\n",
    "            train_loader = data_utils.DataLoader(ColoredMnistBags3_binary(\n",
    "                                          mean_bag_length=bag_size,\n",
    "                                          var_bag_length=0,\n",
    "                                          num_bag=40000//bag_size,\n",
    "                                          seed=1,\n",
    "                                          train=True),\n",
    "                                          batch_size=batch_size, pin_memory=True,num_workers=0,\n",
    "                                          shuffle=True,collate_fn=mi_collate_img)\n",
    "            path = '/home/weights/C_binary_bagsize_'+str(bag_size) + '.pt'\n",
    "        else:\n",
    "            train_loader = data_utils.DataLoader(ColoredFashionMnistBags3_binary(\n",
    "                                          mean_bag_length=bag_size,\n",
    "                                          var_bag_length=0,\n",
    "                                          num_bag=40000//bag_size,\n",
    "                                          train=True),\n",
    "                                          batch_size=batch_size, pin_memory=True,num_workers=0,\n",
    "                                          shuffle=True,collate_fn=mi_collate_img)\n",
    "            path = '/home/weights/C_Fashion_binary_bagsize_'+str(bag_size) + '.pt'\n",
    "\n",
    "        model = training_procedure(FLAGS, (1,28,28), train_loader)\n",
    "        torch.save(model.state_dict(), path)\n",
    "        model.apply(weight_reset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff88d893-93d6-489d-bcff-242605d8673e",
   "metadata": {},
   "source": [
    "#### Evaluation all bags at once"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee7609b7-fa48-403a-a8bc-ab4697ec282c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import roc_auc_score, average_precision_score, precision_recall_fscore_support, accuracy_score, roc_curve\n",
    "\n",
    "def get_accuracy(model, bags, bag_idx, bag_label, instance_label, thresh=0.5):\n",
    "    with torch.no_grad():\n",
    "        pred_instance = model.classifier_ins(bags, bag_idx)\n",
    "    instance_auc = roc_auc_score(instance_label.cpu(), pred_instance.cpu())\n",
    "    instance_aucpr = average_precision_score(instance_label.cpu(), pred_instance.cpu())\n",
    "    \n",
    "    temp = [m > thresh for m in pred_instance.cpu().squeeze().numpy()]\n",
    "    accuracy = accuracy_score(instance_label.cpu(), temp)\n",
    "    precision, recall, fscore, _ = precision_recall_fscore_support(instance_label.cpu(), temp)\n",
    "    return instance_auc, instance_aucpr, accuracy, precision, recall, fscore\n",
    "\n",
    "auc_list = []\n",
    "aucpr_list = []\n",
    "acc_list = []\n",
    "precision_list = []\n",
    "recall_list = []\n",
    "f1_list = []\n",
    "\n",
    "bag_size_test = 2\n",
    "# task = 'mnist'\n",
    "print(task)\n",
    "if task == 'mnist':\n",
    "    test_loader = data_utils.DataLoader(ColoredMnistBags3_binary(\n",
    "                                  mean_bag_length=bag_size_test,\n",
    "                                  var_bag_length=0,\n",
    "                                  num_bag=20000//bag_size_test,\n",
    "                                  train=False),\n",
    "                                  batch_size=20000//bag_size_test, pin_memory=True,num_workers=0,\n",
    "                                  shuffle=True,collate_fn=mi_collate_img)\n",
    "    path = '/home/weights/C_binary_bagsize_'+str(bag_size) + '.pt'\n",
    "else:\n",
    "    test_loader = data_utils.DataLoader(ColoredFashionMnistBags3_binary(\n",
    "                                  mean_bag_length=bag_size_test,\n",
    "                                  var_bag_length=0,\n",
    "                                  num_bag=20000//bag_size_test,\n",
    "                                  train=False), # use test set\n",
    "                                  batch_size=20000//bag_size,\n",
    "                                  shuffle=True,collate_fn=mi_collate_img)\n",
    "    path = '/home/weights/C_Fashion_binary_bagsize_'+str(bag_size) + '.pt'\n",
    "\n",
    "test_bag, test_bag_idx, test_bag_label, test_instance_label = iter(test_loader).next()\n",
    "\n",
    "print(\"number of instances\", len(test_instance_label))\n",
    "print(\"number of positive instances\", sum(test_instance_label))\n",
    "print(\"number of positive bags\", sum(test_bag_label))\n",
    "\n",
    "\n",
    "model = cmil_mnist(FLAGS).to(torch.device('cuda'))\n",
    "model.load_state_dict(torch.load(path))\n",
    "device = torch.device('cuda')\n",
    "\n",
    "test_auc, test_aucpr, accuracy, precision, recall, fscore = get_accuracy(model, test_bag.float().to(device), test_bag_idx, test_bag_label, test_instance_label)\n",
    "print(test_auc)\n",
    "print(test_aucpr)\n",
    "print('Accuracy', accuracy)\n",
    "print('precision', precision)\n",
    "print('recall', recall)\n",
    "\n",
    "auc_list.append(test_auc)\n",
    "aucpr_list.append(test_aucpr)\n",
    "acc_list.append(accuracy)\n",
    "precision_list.append(precision)\n",
    "recall_list.append(recall)\n",
    "f1_list.append(fscore)\n",
    "\n",
    "print('Mean instance prediction AUC is： {:.5f}, AUC-PR: {:.5f}, ACC: {:.5f}'.format(np.mean(auc_list), np.mean(aucpr_list), np.mean(acc_list) ))\n",
    "print('Mean instance prediction Precision is： {:.5f}, Recall: {:.5f}, Fscore: {:.5f}'.format(np.mean(precision_list), np.mean(recall_list), np.mean(f1_list) ))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b55a79c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataloaders.dataloader_draw import ColoredMnistBagsDraw, ColoredFashionMnistBagsDraw_binary\n",
    "\n",
    "def show_img(img):\n",
    "    plt.figure(figsize=(12,12))\n",
    "    npimg = img.numpy()\n",
    "    npimg = np.clip(npimg, 0., 1.)\n",
    "    plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "\n",
    "print(task)\n",
    "bag_size_test = 1\n",
    "i=0\n",
    "if task == 'mnist':\n",
    "    test_loader = data_utils.DataLoader(ColoredMnistBagsDraw(\n",
    "                                  mean_bag_length=bag_size_test,\n",
    "                                  var_bag_length=0,\n",
    "                                  num_bag=20000//bag_size_test,\n",
    "                                  seed=1,\n",
    "                                  train=True),\n",
    "                                  batch_size=batch_size, pin_memory=True,num_workers=0,\n",
    "                                  shuffle=True,collate_fn=mi_collate_img)\n",
    "    path = '/home/user/Code/weights/C_binary2_bagsize_'+str(bag_size)+'reps_'+str(i) + '.pt'\n",
    "else:\n",
    "    test_loader = data_utils.DataLoader(ColoredFashionMnistBagsDraw_binary(\n",
    "                                  mean_bag_length=bag_size_test,\n",
    "                                  var_bag_length=0,\n",
    "                                  num_bag=20000//bag_size_test,\n",
    "                                  train=False), # use test set\n",
    "                                  batch_size=20000//bag_size,\n",
    "                                  shuffle=True,collate_fn=mi_collate_img)\n",
    "    path = '/home/user/Code/weights/C_Fashion_binary_bagsize_'+str(bag_size)+'reps_'+str(i)  + '.pt'\n",
    "\n",
    "model = tmil_mnist(FLAGS).to(torch.device('cuda'))\n",
    "model.load_state_dict(torch.load(path))\n",
    "\n",
    "temp3 = iter(test_loader)\n",
    "batch3 = next(temp3)\n",
    "temp3 = model.get_encoding(batch3[0].to(torch.device('cuda')),batch3[1].to(torch.device('cuda')))\n",
    "img3 = model.reconstruct(temp3)\n",
    "\n",
    "show_img(torchvision.utils.make_grid(batch3[0][0:100], nrow = 10, ncol=10))\n",
    "show_img(torchvision.utils.make_grid(img3[0:100].cpu(), nrow = 10, ncol=10))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.11 ('pytorch')",
   "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.11"
  },
  "vscode": {
   "interpreter": {
    "hash": "5839d9b809a0d3e5138897fed927252c680de6913018283041a74cef55b79d30"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}