mace+maxpool_visualize.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "TL3VmtHIcr0Q",
    "outputId": "22389230-971b-4f08-a43e-289e6d4a5450"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "from torch.utils.data import Dataset,DataLoader\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "from torchvision import transforms, datasets\n",
    "import time\n",
    "import os\n",
    "import copy\n",
    "import gc\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "import PIL\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "ebJp0v9dFMKt"
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "EY6w3qwDqTIu"
   },
   "outputs": [],
   "source": [
    "class EmbeddingGenerator(nn.Module):\n",
    "    def __init__(self, activation_features_channels=200, activation_features_size=(14, 14)):\n",
    "        \"\"\"     \n",
    "        Parameters\n",
    "        ----------\n",
    "        activation_features_channels : INT, optional            DESCRIPTION. The default is 512 as in VGG.\n",
    "        activation_features_size : TYPE, optional            DESCRIPTION. The default is (14, 14) as in VGG.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        None.\n",
    "        \"\"\"\n",
    "        super().__init__()             \n",
    "        self.embeddings = nn.Sequential( nn.Linear(activation_features_size[0] * activation_features_size[1], 128),\n",
    "                                        nn.ReLU(),\n",
    "                                        nn.Linear(128, 64),\n",
    "                                        nn.ReLU(),\n",
    "                                        nn.Linear(64, 32),\n",
    "                                        nn.Tanh(),)\n",
    "        self.embeddings = nn.Sequential( nn.Linear(activation_features_size[0] * activation_features_size[1], 32),\n",
    "                                        nn.Tanh(),)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \"\"\" input will be in concept map dimensions, so first flatten it then pass through embedding generator \"\"\"        \n",
    "        x = torch.flatten(x, 1)        \n",
    "        x = self.embeddings(x)\n",
    "        weights = torch.sum(torch.square(x),dim=-1,keepdim=True)       \n",
    "        x = F.normalize(x)        \n",
    "        return x, weights   \n",
    "    \n",
    "class Relevance(nn.Module):\n",
    "    def __init__(self, proto_per_class = 10,embed_out_shape=(32,32)):\n",
    "        super(Relevance , self).__init__()\n",
    "        self.rel = nn.Linear(embed_out_shape[0] * proto_per_class, 1)\n",
    "\n",
    "    def forward(self,inputs):\n",
    "        relevance = self.rel(inputs)\n",
    "        return relevance     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "FsmqoZBiZthC"
   },
   "outputs": [],
   "source": [
    "class MACENetwork(nn.Module):\n",
    "    \"\"\" MACE module \"\"\"\n",
    "    def __init__(self, activation_features_channels = 512 , num_classes = 20, proto_per_class = 10, first_dense_dim = 4096):\n",
    "        super(MACENetwork , self).__init__()\n",
    "\n",
    "        self.activation_features_channels = activation_features_channels\n",
    "        self.num_classes = num_classes\n",
    "        self.proto_per_class = proto_per_class\n",
    "        self.num_prototypes = num_classes * proto_per_class\n",
    "\n",
    "        # concept map generator module\n",
    "        self.feature_extract = nn.Sequential(nn.Conv2d(in_channels=self.activation_features_channels,out_channels=self.num_prototypes,\n",
    "                                                       kernel_size=(1,1)), nn.ReLU())\n",
    "    \n",
    "        # maps a concept map into a local prototype(embedding generator)\n",
    "        self.protonet = nn.ModuleList([EmbeddingGenerator(activation_features_size=(1, 1)) for _ in range(self.num_classes)])\n",
    "    \n",
    "        # concept relevance module\n",
    "        self.concept_relevance  = nn.ModuleList([Relevance(self.proto_per_class) for _ in range(self.num_classes)])\n",
    "\n",
    "        self.softmax = nn.Sigmoid()\n",
    "        \n",
    "        # Final Dense Layer (output module)\n",
    "        self.final_dense = nn.Sequential(nn.Linear(32*self.num_prototypes , first_dense_dim), nn.ReLU())\n",
    "\n",
    "    def get_relevance_part(self):\n",
    "        return self.concept_relevance\n",
    "    \n",
    "    def forward(self, inputs):\n",
    "        batch_size = inputs.shape[0]\n",
    "        # Get concept maps\n",
    "        concept_maps = self.feature_extract(inputs)\n",
    "         # for maxpool\n",
    "        concept_maps = F.max_pool2d(concept_maps, kernel_size=(concept_maps.size()[2], concept_maps.size()[3]))\n",
    "        concept_maps = torch.unsqueeze(concept_maps, -1)\n",
    "        concept_maps = torch.unsqueeze(concept_maps, -1)\n",
    "        concept_maps = torch.unbind(concept_maps, dim = 1) \n",
    "        \n",
    "        # List of lists for prototypes(or embeddings)\n",
    "        class_prototypes = [[] for _ in range(self.num_classes)]\n",
    "        class_protoweight = [[] for _ in range(self.num_classes)]\n",
    "        class_protoweighted = [[] for _ in range(self.num_classes)]\n",
    "        \n",
    "        # Get Concepts\n",
    "        for i, concept_map in enumerate(concept_maps):\n",
    "            # class index ci\n",
    "            ci = i//self.proto_per_class\n",
    "            \n",
    "            # Pass the concept map to the protonet\n",
    "            prototype, weight = self.protonet[ci](concept_map)\n",
    "            class_prototypes[ci].append(prototype)\n",
    "            class_protoweight[ci].append(weight)\n",
    "\n",
    "            #if self.prune[ci][i % self.proto_per_class]:\n",
    "            #    prototype = 0 * prototype\n",
    "            class_protoweighted[ci].append(prototype)\n",
    "        \n",
    "        # relevances and concepts\n",
    "        relevances, concepts = [], []\n",
    "\n",
    "        # Get Relevances\n",
    "        for i, cprototypes in enumerate(class_protoweighted):\n",
    "            # class prototypes to tensors      # (batch_size, embedding_size * prototypes per class)      \n",
    "            cps = torch.cat(cprototypes, 1)             \n",
    "            relev = self.concept_relevance[i](cps)            \n",
    "\n",
    "            # Add Relevances and concepts\n",
    "            relevances.append(relev)\n",
    "            concepts.append(cps);\n",
    "        \n",
    "        # (Batch_size, num_classes)\n",
    "        relevances = torch.cat(relevances,1)   \n",
    "        relevances = self.softmax(relevances)        \n",
    "        expanded_relevances = torch.unsqueeze(relevances,-1)  \n",
    "            \n",
    "        # (Batch_size, embedding_size * num prototypes)\n",
    "        concepts = torch.cat(concepts,1) \n",
    "       \n",
    "        temp = torch.ones((batch_size , self.num_classes , 32*self.proto_per_class),requires_grad=False)\n",
    "        temp = temp.to(device)\n",
    "        expanded_relevances = expanded_relevances * temp  \n",
    "        expanded_relevances = torch.reshape(expanded_relevances, concepts.shape)\n",
    "        \n",
    "        # Weight the concepts\n",
    "        weighted_concepts = expanded_relevances * concepts\n",
    "        dense_layer = self.final_dense(weighted_concepts)\n",
    "        \n",
    "        return class_prototypes, class_protoweight, relevances, dense_layer, class_protoweighted   \n",
    "\n",
    "    def get_vis_local(self,inputs):\n",
    "        batch_size = inputs.shape[0]\n",
    "        concept_maps = self.feature_extract(inputs)\n",
    "        concept_maps_to_return = torch.unbind(concept_maps, dim = 1)\n",
    "        concept_maps = F.max_pool2d(concept_maps, kernel_size=(concept_maps.size()[2], concept_maps.size()[3]))\n",
    "        concept_maps = torch.unsqueeze(concept_maps, -1)\n",
    "        concept_maps = torch.unsqueeze(concept_maps, -1)\n",
    "        concept_maps = torch.unbind(concept_maps, dim = 1)    \n",
    "\n",
    "        # List of lists for prototypes\n",
    "        class_prototypes = [[] for i in range(self.num_classes)]\n",
    "        class_protoweight = [[] for i in range(self.num_classes)]\n",
    "        class_protoweighted = [[] for i in range(self.num_classes)]\n",
    "        \n",
    "        # Get Concepts\n",
    "        for i, concept_map in enumerate(concept_maps):\n",
    "            # class index ci\n",
    "            ci = i//self.proto_per_class        \n",
    "            # Pass the featuremap to the protonet\n",
    "            prototype, weight = self.protonet[ci](concept_map)\n",
    "            class_prototypes[ci].append(prototype)\n",
    "            class_protoweight[ci].append(weight)\n",
    "\n",
    "            #if self.prune[ci][i % self.proto_per_class]:\n",
    "            #    prototype = 0 * prototype\n",
    "            class_protoweighted[ci].append(prototype)\n",
    "\n",
    "        # relevances and concepts\n",
    "        relevances, concepts = [], []\n",
    "        \n",
    "        # Get Relevances\n",
    "        for i, cprototypes in enumerate(class_protoweighted):\n",
    "            cps = torch.cat(cprototypes, 1)           # class prototypes to tensors (batch_size, embedding_size * prototypes per class)\n",
    "            relev = self.concept_relevance[i](cps)\n",
    "        \n",
    "            # Add Relevances and concepts\n",
    "            relevances.append(relev)\n",
    "            concepts.append(cps);\n",
    "\n",
    "        # (Batch_size, num_classes)\n",
    "        relevances = torch.cat(relevances,1)    \n",
    "        relevances = self.softmax(relevances)\n",
    "     \n",
    "        return concept_maps_to_return, class_prototypes, class_protoweight, relevances, class_protoweighted            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "CZEIFgRkFMKt"
   },
   "outputs": [],
   "source": [
    "class ApplyMACE(nn.Module):\n",
    "    def __init__(self, activation_features_channels = 512,num_classes = 20, proto_per_class = 10, first_dense_dim = 4096):\n",
    "        super().__init__()\n",
    "        self.basemodel = load_checkpoint('sgd_vgg16_finetune_checkpoint.pth')\n",
    "        self.interpret_layer = MACENetwork(activation_features_channels,num_classes, proto_per_class, first_dense_dim)\n",
    "        \n",
    "        self.layer_outputs = {}\n",
    "        def get_activation(name):\n",
    "            def hook(module, input, output):\n",
    "                self.layer_outputs[name] = output\n",
    "            return hook\n",
    "        self.basemodel.features[30].register_forward_hook(get_activation('convolution_output'))\n",
    "        self.basemodel.classifier[1].register_forward_hook(get_activation('first_fully_connected_layer_output'))\n",
    "        \n",
    "        for param in self.basemodel.parameters():\n",
    "            param.requires_grad = False\n",
    "        \n",
    "        for params in self.interpret_layer.parameters():\n",
    "            params.requires_grad = True\n",
    "            \n",
    "    def get_features(self, inputs):\n",
    "        final_output = self.basemodel(inputs)\n",
    "        final_output = F.softmax(final_output,dim=1)\n",
    "        features =  self.layer_outputs['convolution_output']\n",
    "        first_dense = self.layer_outputs['first_fully_connected_layer_output']\n",
    "        return features, first_dense, final_output\n",
    "    \n",
    "    #get_softmax only for vgg model\n",
    "    def get_softmax(self, inputs):\n",
    "        for i in range(3,7):\n",
    "            inputs=self.basemodel.classifier[i](inputs)    \n",
    "        result = F.softmax(inputs,dim=1) \n",
    "        #result = inputs\n",
    "        return result\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        features, first_dense, final_output = self.get_features(inputs)\n",
    "        class_prototypes, class_protoweight, relevances, dense_layer_predict, class_protoweighted = self.interpret_layer(features)\n",
    "        final_predict = self.get_softmax(dense_layer_predict)\n",
    "        return class_prototypes, class_protoweight, relevances, dense_layer_predict, first_dense, final_predict, final_output, class_protoweighted\n",
    "    \n",
    "    def get_concept_maps(self, inputs):\n",
    "        features, first_dense, final_output = self.get_features(inputs)\n",
    "        return self.interpret_layer.get_concept_maps(features)\n",
    "    \n",
    "    def get_vis_local_maps(self,inputs):\n",
    "        features, _, _ = self.get_features(inputs)\n",
    "        return self.interpret_layer.get_vis_local(features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "eLQHs-1ACy4_"
   },
   "outputs": [],
   "source": [
    "def load_checkpoint(filepath):\n",
    "    checkpoint = torch.load(filepath)\n",
    "    model = checkpoint['model']\n",
    "    model.load_state_dict(checkpoint['state_dict'])\n",
    "    for parameter in model.parameters():\n",
    "        parameter.requires_grad = False\n",
    "    \n",
    "    model.eval()\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "meta_model = ApplyMACE()\n",
    "meta_model.to(device)\n",
    "\n",
    "ch_pt = torch.load('./max_pool_92acc.pth')\n",
    "meta_model.load_state_dict(ch_pt['model_state_dict'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "20"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_dataset_loader 57\n"
     ]
    }
   ],
   "source": [
    "data_transform = transforms.Compose([transforms.Resize((224,224)),\n",
    "                                     transforms.ToTensor(),\n",
    "                                     transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])])\n",
    "\n",
    "train_dataset = datasets.ImageFolder(root='./DATA/output/splitted_train_data/val', transform=data_transform)\n",
    "train_dataset_loader = torch.utils.data.DataLoader(train_dataset, batch_size = 64, shuffle=True, num_workers=2)\n",
    "\n",
    "print('train_dataset_loader',len(train_dataset_loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 3, 224, 224])\n",
      "57\n",
      "5\n",
      "20\n",
      "10\n",
      "torch.Size([64, 32])\n",
      "torch.Size([64, 32])\n",
      "torch.Size([64, 32])\n",
      "torch.Size([64, 32])\n",
      "torch.Size([64, 32])\n",
      "torch.Size([64, 32])\n",
      "torch.Size([64, 32])\n",
      "torch.Size([64, 32])\n",
      "torch.Size([64, 32])\n",
      "torch.Size([64, 32])\n",
      "torch.Size([64, 32])\n",
      "concept_map torch.Size([200, 3592, 7, 7])\n",
      "class_prototype torch.Size([20, 10, 3592, 32])\n",
      "class_protoweight torch.Size([20, 10, 3592, 1])\n",
      "class_protoweight torch.Size([20, 10, 3592])\n",
      "relevance torch.Size([3592, 20])\n"
     ]
    }
   ],
   "source": [
    "def get_maps(train_set_ims):\n",
    "    all_outputs = [meta_model.get_vis_local_maps(image.to(device)) for image, label in train_set_ims]\n",
    "    print(len(all_outputs))\n",
    "    print(len(all_outputs[0]))\n",
    "    print(len(all_outputs[0][1]))\n",
    "    print(len(all_outputs[0][1][0]))\n",
    "    print((all_outputs[0][1][0][0].shape))\n",
    "    for i in all_outputs[0][1][0]:\n",
    "        print(i.shape)\n",
    "   \n",
    "    concept_map = torch.cat([torch.stack(x[0]) for x in all_outputs], 1)\n",
    "    print('concept_map',concept_map.size())\n",
    "    \n",
    "    class_prototype = torch.cat([torch.stack([torch.stack(i) for i in x[1]]) for x in all_outputs], 2)\n",
    "    print('class_prototype',class_prototype.shape)\n",
    "    \n",
    "    class_protoweight = torch.cat([torch.stack([torch.stack(i) for i in x[2]]) for x in all_outputs], 2)\n",
    "    print('class_protoweight',class_protoweight.shape)\n",
    "    \n",
    "    class_protoweight = torch.squeeze(class_protoweight)\n",
    "    print('class_protoweight',class_protoweight.shape)\n",
    "    \n",
    "    relevance = torch.cat([x[3] for x in all_outputs])\n",
    "    print('relevance',relevance.shape)\n",
    "    return concept_map, class_prototype, class_protoweight, relevance\n",
    "\n",
    "concept_maps, class_prototypes, class_protoweights, relevances = get_maps(train_dataset_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "DIRECTORY = 'project'\n",
    "dir_v = 'visualizations_maxpool_after_clip/'\n",
    "if not os.path.exists(dir_v):\n",
    "    os.makedirs(dir_v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_weights(layer_weights, prototype, i ,j):\n",
    "    # print(layer_weights[i].rel)\n",
    "    weights = layer_weights[i].rel.weight[0]\n",
    "    # print(weights.shape)\n",
    "    a = torch.split(weights, [32 for _ in range(10)], 0)\n",
    "    # print(a[j].shape)\n",
    "    # print('torch.matmul(prototype, a[j])',torch.matmul(prototype, a[j]).shape)\n",
    "    return torch.matmul(prototype, a[j]).squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize(num_classes=20, proto_per_class=10, topK=10,S=5): # train_dataset_loader\n",
    "    mean = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32)\n",
    "    std = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32)\n",
    "    denormalize = transforms.Normalize((-mean / std).tolist(), (1.0 / std).tolist())\n",
    "    compose = transforms.Compose([transforms.Resize((224,224))])\n",
    "    trans = transforms.ToPILImage()\n",
    "    layer_weights = meta_model.interpret_layer.get_relevance_part()\n",
    "   \n",
    "    for i in range(num_classes):\n",
    "        for j in range(proto_per_class):\n",
    "            #print('concept_maps',concept_maps.shape)  torch.Size([200, 3592, 7, 7])\n",
    "            c_m = concept_maps[i*10 + j]\n",
    "            \n",
    "            c_pw = class_protoweights[i][j]\n",
    "            c_prototype = class_prototypes[i][j]\n",
    "            new_c_pw = get_weights(layer_weights, c_prototype, i ,j)\n",
    "            weight_rel = new_c_pw\n",
    "            \n",
    "            #values,indices = torch.topk( weight_rel , dim=0, k = topK )\n",
    "            indices = torch.nonzero(torch.eq(torch.tensor(train_dataset.targets),i).float()).flatten()\n",
    "\n",
    "            cm = [c_m[index]  for index in indices]\n",
    "            names = [train_dataset[x][0] for x in indices]\n",
    "            targets = sum([1 for _i in [train_dataset[x][1] for x in indices] if _i==i])\n",
    "            if targets < 5:\n",
    "                continue\n",
    "                \n",
    "            for k in range(topK):\n",
    "                image = names[k].to(device)\n",
    "                cmi = cm[k]\n",
    "                #print(cmi)\n",
    "                cmi = cmi.eq(torch.max(cmi)).float()\n",
    "                #print(cmi)\n",
    "                cmi = torch.clamp(cmi,0.38,1)\n",
    "                #print(cmi)\n",
    "                \n",
    "                cmi = torch.unsqueeze(cmi, 0)                \n",
    "                map_img = transforms.Resize((224, 224), interpolation=PIL.Image.NEAREST)(cmi)\n",
    "                image = denormalize(image)\n",
    "                image = torch.clamp(image, 0, 1)\n",
    "\n",
    "                map_img = map_img*image\n",
    "                # map_img = map_img/torch.max(map_img)\n",
    "                \n",
    "                map_name = dir_v + '/class' + str(i) + '_concept' + str(j) + '_' + str(k) + '_m_' + '.png'\n",
    "                im_name  = dir_v + '/class' + str(i) + '_concept' + str(j) + '_' + str(k) + '_i_' + '.png'\n",
    "                \n",
    "                torchvision.utils.save_image(map_img, map_name)\n",
    "                torchvision.utils.save_image(image, im_name)                \n",
    "\n",
    "visualize(num_classes=20, topK=10)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "ann_final.ipynb",
   "provenance": []
  },
  "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}