{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffa43315-fd5b-4189-8195-29aab58174b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "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 numpy as np\n",
    "import torch\n",
    "import torch.utils.data as data_utils\n",
    "from torchvision import datasets, transforms\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\n",
    "\n",
    "\n",
    "from dataloaders.dataloader import KMnistBags3, MnistBags3, FashionMnistBags3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1ee5800-9717-4521-a017-d9802992e678",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79fab365-eec9-4483-9890-d36dbbd27e55",
   "metadata": {},
   "outputs": [],
   "source": [
    "from util import *\n",
    "from models.networks_mnists import cmil_mnist"
   ]
  },
  {
   "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 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",
    "    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(auto_encoder_optimizer, factor=0.1, patience=20, verbose=True)\n",
    "\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\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",
    "\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",
    "      scheduler.step(y_epoch)\n",
    "\n",
    "      if ((epoch + 1) % 10 ==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",
    "        \n",
    "      if elbo_epoch < best_loss:\n",
    "            best_loss = elbo_epoch\n",
    "            torch.save(model.state_dict(), '/home/user/temp_'+ FLAGS.task +'.pt')\n",
    "\n",
    "    model.load_state_dict(torch.load('/home/user/temp_'+ FLAGS.task +'.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",
    "    \n",
    "\n",
    "#     task = 'fashion'\n",
    "    # task = 'mnist'\n",
    "    task = 'kmnist'\n",
    "    \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 = 1, help=\"input channels\")\n",
    "        parser.add_argument('--task', type=str, default = task, help=\"current task\")\n",
    "\n",
    "        FLAGS = parser.parse_args(args=[])\n",
    "\n",
    "        batch_size = 2048\n",
    "        bag_size = 30\n",
    "        \n",
    "\n",
    "        n_tasks = 10\n",
    "        print('Current Task is', task)\n",
    "        \n",
    "        for target in range(0, n_tasks):\n",
    "            print(target)\n",
    "            if task == 'mnist':\n",
    "                train_loader = data_utils.DataLoader(MnistBags3(target_number=target,\n",
    "                                              mean_bag_length=bag_size,\n",
    "                                              var_bag_length=0,\n",
    "                                              num_bag=60000//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/user/Code/weights/set_mnist_'+str(target) + '.pt'\n",
    "            elif task =='fashion':\n",
    "                train_loader = data_utils.DataLoader(FashionMnistBags3(target_number=target,\n",
    "                                              mean_bag_length=bag_size,\n",
    "                                              var_bag_length=0,\n",
    "                                              num_bag=60000//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/user/Code/weights/set_fmnist_'+str(target) + '.pt'\n",
    "            elif task =='kmnist':\n",
    "                train_loader = data_utils.DataLoader(KMnistBags3(target_number=target,\n",
    "                                              mean_bag_length=bag_size,\n",
    "                                              var_bag_length=0,\n",
    "                                              num_bag=60000//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/user/Code/weights/set_kmnist_'+str(target) + '.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",
    "import sklearn\n",
    "\n",
    "def get_accuracy_multiclass(model, bags, bag_idx, bag_label, instance_label, threshold=0.5):\n",
    "    with torch.no_grad():\n",
    "        pred_instance = model.classifier_ins(bags, bag_idx)\n",
    "    temp = [m > threshold for m in pred_instance.cpu().squeeze().numpy()]\n",
    "    \n",
    "    precision = sklearn.metrics.precision_score(instance_label, temp) # tp/ tp+fp\n",
    "    recall = sklearn.metrics.recall_score(instance_label, temp) # tp/ tp + fn\n",
    "    \n",
    "    tp = precision * sum(temp)\n",
    "    print('tp', tp)\n",
    "    fp = sum(temp) - tp\n",
    "    print('fp',fp)\n",
    "    \n",
    "    \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",
    "    accuracy = accuracy_score(instance_label.cpu(), temp)\n",
    "    precision, recall, fscore, _ = precision_recall_fscore_support(instance_label.cpu(), temp, average='binary')\n",
    "    return instance_auc, instance_aucpr, accuracy, precision, recall, fscore, tp, fp\n",
    "\n",
    "auc_list = []\n",
    "aucpr_list = []\n",
    "acc_list = []\n",
    "precision_list = []\n",
    "recall_list = []\n",
    "f1_list = []\n",
    "tp_list = []\n",
    "fp_list = []\n",
    "\n",
    "# task = 'fashion'\n",
    "# task = 'mnist'\n",
    "# task = 'kmnist'\n",
    "print('Current Task is', task)\n",
    "for target in range(0,10):\n",
    "    print(target)\n",
    "    if task == 'mnist':\n",
    "        test_loader = data_utils.DataLoader(MnistBags3(target_number=target,\n",
    "                                      mean_bag_length=bag_size,\n",
    "                                      var_bag_length=0,\n",
    "                                      num_bag=60000//bag_size,\n",
    "                                      seed=1,\n",
    "                                      train=False),\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/set_mnist_'+str(target) + '.pt'\n",
    "    elif task =='fashion':\n",
    "        test_loader = data_utils.DataLoader(FashionMnistBags3(target_number=target,\n",
    "                                      mean_bag_length=bag_size,\n",
    "                                      var_bag_length=0,\n",
    "                                      num_bag=60000//bag_size,\n",
    "                                      train=False),\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/set_fmnist_'+str(target) + '.pt'\n",
    "    elif task =='kmnist':\n",
    "        test_loader = data_utils.DataLoader(KMnistBags3(target_number=target,\n",
    "                                      mean_bag_length=bag_size,\n",
    "                                      var_bag_length=0,\n",
    "                                      num_bag=60000//bag_size,\n",
    "                                      train=False),\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/set_kmnist_'+str(target) + '.pt'\n",
    "        \n",
    "    test_bag, test_bag_idx, test_bag_label, test_instance_label = iter(test_loader).next()\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,tp, fp = get_accuracy_multiclass(model, test_bag.float().to(device), \n",
    "                                                                           test_bag_idx, test_bag_label, test_instance_label)\n",
    "    print('Target is', target)\n",
    "    print('AUC', test_auc)\n",
    "    print('MAP', 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",
    "    tp_list.append(tp)\n",
    "    fp_list.append(fp)\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",
    "\n",
    "print('Mean instance prediction Accuracy: {:.5f}'.format(np.sum(tp_list)/(np.sum(tp_list)+np.sum(fp_list)) ))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66583121-f758-4c29-bee8-39b5d377191f",
   "metadata": {},
   "source": [
    "#### Qualitative Evaluation Random Order"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d45d24c8-1163-4712-8adb-5e2ce929edc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_img(img):\n",
    "    plt.figure(figsize=(12,10))\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",
    "# task = 'fashion'\n",
    "#         task = 'mnist'\n",
    "#         task = 'kmnist'\n",
    "print('Current Task is', task)\n",
    "for target in range(0,n_tasks):\n",
    "    print(target)\n",
    "    if task == 'mnist':\n",
    "        test_loader = data_utils.DataLoader(MnistBags3(target_number=target,\n",
    "                                      mean_bag_length=bag_size,\n",
    "                                      var_bag_length=0,\n",
    "                                      num_bag=60000//bag_size,\n",
    "                                      train=False),\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/set_mnist_'+str(target) + '.pt'\n",
    "    elif task =='fashion':\n",
    "        test_loader = data_utils.DataLoader(FashionMnistBags3(target_number=target,\n",
    "                                      mean_bag_length=bag_size,\n",
    "                                      var_bag_length=0,\n",
    "                                      num_bag=60000//bag_size,\n",
    "                                      train=False),\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/set_fmnist_'+str(target) + '.pt'\n",
    "    elif task =='kmnist':\n",
    "        test_loader = data_utils.DataLoader(KMnistBags3(target_number=target,\n",
    "                                      mean_bag_length=bag_size,\n",
    "                                      var_bag_length=0,\n",
    "                                      num_bag=60000//bag_size,\n",
    "                                      train=False),\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/set_kmnist_'+str(target) + '.pt'\n",
    "        \n",
    "    \n",
    "    temp2 = iter(test_loader)\n",
    "    batch2= next(temp2)\n",
    "\n",
    "    model = cmil_mnist(FLAGS).to(torch.device('cuda'))\n",
    "    model.load_state_dict(torch.load(path))\n",
    "\n",
    "\n",
    "    # show_img(torchvision.utils.make_grid(img[0:100].cpu(), nrow=10, ncol=10))\n",
    "    # show_img(torchvision.utils.make_grid(batch[0][0:100], nrow=10, ncol=10))\n",
    "\n",
    "\n",
    "    temp2 = model.get_encoding(batch2[0].to(torch.device('cuda')),batch2[1].to(torch.device('cuda')))\n",
    "    img2 = model.reconstruct(temp2)\n",
    "\n",
    "    print(temp2.shape)\n",
    "    print(img2.shape)\n",
    "\n",
    "    show_img(torchvision.utils.make_grid(img2[0:100].cpu(), nrow=10, ncol=10))\n",
    "    show_img(torchvision.utils.make_grid(batch2[0][0:100], nrow=10, ncol=10))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8ea7356-d7d5-41c3-b1b4-6371a57582d7",
   "metadata": {},
   "source": [
    "### Qualitative Ordered"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c4a8d98-6b5c-4d86-a82b-a81a14f09e19",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataloaders.dataloader_draw import KMnistBagsDraw, MnistBagsDraw, FashionMnistBagsDraw\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('Current Task is', task)\n",
    "for target in range(0,10):\n",
    "    print(target)\n",
    "    if task == 'mnist':\n",
    "        test_loader = data_utils.DataLoader(MnistBagsDraw(target_number=target,\n",
    "                                      mean_bag_length=bag_size,\n",
    "                                      var_bag_length=0,\n",
    "                                      num_bag=60000//bag_size,\n",
    "                                      seed=1,\n",
    "                                      train=False),\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/set_mnist_'+str(target) + '.pt'\n",
    "    elif task =='fashion':\n",
    "        test_loader = data_utils.DataLoader(FashionMnistBagsDraw(target_number=target,\n",
    "                                      mean_bag_length=bag_size,\n",
    "                                      var_bag_length=0,\n",
    "                                      num_bag=60000//bag_size,\n",
    "                                      train=False),\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/set_fmnist_'+str(target) + '.pt'\n",
    "    elif task =='kmnist':\n",
    "        test_loader = data_utils.DataLoader(KMnistBagsDraw(target_number=target,\n",
    "                                      mean_bag_length=bag_size,\n",
    "                                      var_bag_length=0,\n",
    "                                      num_bag=60000//bag_size,\n",
    "                                      train=False),\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/set_kmnist_'+str(target) + '.pt'\n",
    "        \n",
    "\n",
    "    model = cmil_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
}
