mace+maxpool.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"
   ]
  },
  {
   "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",
    "        \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",
    "\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_to_return, dim = 1)\n",
    "        \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",
    "        \n",
    "        concept_maps = torch.unbind(concept_maps, dim = 1)   \n",
    "         \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",
    "        # print(concept_maps[0].shape)\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": 7,
   "metadata": {
    "id": "FAdBKhQzlDKI"
   },
   "outputs": [],
   "source": [
    "def get_anchors(batch_size, num_concepts):\n",
    "    anchors_indices = torch.from_numpy(np.random.uniform(low = 0, high = batch_size, size=(num_concepts, )))\n",
    "    anchors_indices = anchors_indices.type(torch.int32)\n",
    "    index_mapper = batch_size * (torch.range(1,num_concepts))\n",
    "    return anchors_indices + index_mapper\n",
    "\n",
    "def pairwise_dist_map(anchors, embeddings):\n",
    "    # RETURNS PAIRWISE DISTANCE FOR EACH CONCEPT\n",
    "    fn = lambda x: torch.sum(torch.square(x - embeddings), axis = 1)\n",
    "    y = torch.stack([fn(ele) for ele in torch.unbind(anchors)])\n",
    "    return y\n",
    "    \n",
    "def sample_negative(avg_positive, masked_negative):\n",
    "    # GIVEN A PANCHOR, IT GETS THE SEMI-HARD NEGATIVE\n",
    "    def get_negative(panchor, nanchors):\n",
    "        mask = torch.gt(nanchors, panchor)                 \n",
    "        applied_map = torch.masked_select(nanchors, mask)   \n",
    "        semi_hard_negative = torch.min(applied_map)          \n",
    "        return semi_hard_negative if torch.isfinite(semi_hard_negative) else torch.tensor(0.0) \n",
    "  \n",
    "    # FOR EACH ELEMENT IN POSITIVE MASK, FIND A SEMI-HARD NEGATIVE\n",
    "    func = lambda x: get_negative(x[0], x[1])\n",
    "    yy = torch.stack([func(ele) for ele in zip(torch.unbind(avg_positive),torch.unbind(masked_negative))])\n",
    "    yy = yy.type(torch.float32)\n",
    "    return yy                                                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "SHFkiy0umL3v"
   },
   "outputs": [],
   "source": [
    "def custom_triplet_loss(local_prototypes, margin = 0.2):\n",
    "    \"\"\" Memory Efficient Triplet Loss \"\"\"\n",
    "    batch_size, embedding_size, num_concepts = local_prototypes.shape\n",
    "\n",
    "    # EMBEDDING: (1st CONCEPT FOR BATCH_SIZE, 2ND CONCEPT FOR BATCH_SIZE)   # SHAPE = (BATCH_SIZE * NUM_CONCEPTS, EMBEDDING_SIZE)\n",
    "    transposed_prototype = local_prototypes.permute(2,0,1)                                            \n",
    "    embeddings = torch.reshape(transposed_prototype, (batch_size*num_concepts, embedding_size))\n",
    "\n",
    "    # MASKS = (NUM_CONCEPTS, NUM_CONCEPT*BATCH_SIZE)\n",
    "    concepts_range = torch.unsqueeze(torch.range(1,num_concepts),-1)                                  \n",
    "    column_mask = torch.tensor(np.broadcast_to(concepts_range.numpy(), (num_concepts, batch_size)))   \n",
    "    column_mask = torch.reshape(column_mask, (num_concepts*batch_size,1))\n",
    "    column_mask = torch.transpose(torch.tensor(np.broadcast_to(column_mask.numpy(),(num_concepts*batch_size, num_concepts))), 0, 1)  \n",
    "\n",
    "    positive_mask = torch.eq(concepts_range, column_mask)       \n",
    "    negative_mask = torch.logical_not(positive_mask)            \n",
    "\n",
    "    # DIMENSION OF ANCHORS = (NUM_CONCEPTS, EMBEDDING_SIZE)\n",
    "    anchor_indices = get_anchors(batch_size, num_concepts)\n",
    "    anchors = torch.gather(embeddings,0, anchor_indices)          \n",
    "\n",
    "    # PAIRWISE DISTANCE = (NUM_CONCEPTS, NUM_CONCEPTS*BATCH_SIZE)\n",
    "    pairwise_distance = pairwise_dist_map(anchors, embeddings)\n",
    "\n",
    "    # FIND HARDEST POSITIVE ANCHOR\n",
    "    masked_positive = torch.masked_select(pairwise_distance, positive_mask)                   \n",
    "    masked_positive = torch.reshape(masked_positive, (num_concepts, batch_size))\n",
    "    avg_positive = torch.mean(masked_positive, dim = 1)                                       \n",
    "\n",
    "    # SEMI HARD NEGATIVE\n",
    "    masked_negative = torch.masked_select(pairwise_distance, negative_mask)                         \n",
    "    masked_negative = torch.reshape(masked_negative, (num_concepts, batch_size*(num_concepts - 1)))\n",
    "    hardest_negative = torch.unsqueeze(sample_negative(avg_positive, masked_negative), -1)          \n",
    "\n",
    "    loss = torch.mean(torch.maximum(masked_positive - hardest_negative + margin, 0))           \n",
    "\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "dqLcso1mz9uq"
   },
   "outputs": [],
   "source": [
    "def new_triplet_loss(prototypes, margin = 0.2):\n",
    "    \"\"\" Memory Efficient Triplet Loss      prototypes.shape = (batch_size, number_of_prototypes, embedding_size) \"\"\"\n",
    "    # FIND HARDEST POSITIVE ACROSS BATCH\n",
    "    p_prototypes = prototypes.permute(1,0,2)                 \n",
    "    p_prototypes = torch.unsqueeze(p_prototypes, -1)         \n",
    "    p_ = p_prototypes\n",
    "    p_t = p_prototypes.permute(0, 3, 2, 1)                   \n",
    "    \n",
    "    pairwise = torch.sum(torch.square(p_ - p_t), dim = 2)    \n",
    "    # Diagonals are 0. For each row take the one with the max.\n",
    "    hardest_p,_ = torch.max(pairwise, dim = -1)               # (number_of_prototypes, batch_size)     \n",
    "    hardest_p = hardest_p.permute(1,0)                       # (batch_size, number_of_prototypes)     \n",
    "    \n",
    "    # FIND HARDEST NEGATIVE ACROSS CONCEPTS\n",
    "    n_prototypes = torch.unsqueeze(prototypes, -1)           \n",
    "    n_ = n_prototypes\n",
    "    n_t = n_prototypes.permute(0, 3, 2, 1)                   \n",
    "    \n",
    "    pairwise = torch.sum(torch.square(n_ - n_t), dim = 2)                      \n",
    "    semi_masked = torch.gt(pairwise, torch.unsqueeze(hardest_p, -1))            \n",
    "    not_semi_masked =  torch.logical_not(semi_masked)                           \n",
    "    # Diagonals are 0. Make diagonals 4 + 1e-6 (largest distance) For each row take the one with the min. \n",
    "    not_semi_masked = not_semi_masked.type(torch.float32)\n",
    "    eye = (4 + 1e-6)*(not_semi_masked)                                         \n",
    "    semi_masked = semi_masked.type(torch.float32)\n",
    "    pairwise = eye + pairwise*semi_masked                                   \n",
    "    \n",
    "    hardest_n , _ = torch.min(pairwise, dim = -1)       # (batch_size, number_of_prototypes)  \n",
    "    temp = torch.zeros(hardest_p.shape,requires_grad = False).to(device)\n",
    "    \n",
    "    x = torch.max(hardest_p - hardest_n + margin,temp)                  #element wise max #(batch_size, number_of_prototypes)\n",
    "    \n",
    "    loss = torch.mean(x, dim = -1)     \n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4dAqVzQZriRG",
    "outputId": "14a2fa59-1239-4207-db90-8c0302b0df4f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_dataset_loader 374\n"
     ]
    }
   ],
   "source": [
    "data_transform = transforms.Compose([transforms.RandomResizedCrop(224),\n",
    "                                     transforms.RandomHorizontalFlip(),\n",
    "                                     transforms.ToTensor(),\n",
    "                                     transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])])\n",
    "train_dataset = datasets.ImageFolder(root='./DATA/output/train', transform=data_transform)\n",
    "#train_dataset = datasets.ImageFolder(root='/content/gdrive/My Drive/iitropar/ANN/P29-Cross-Pure-Dogs/train', transform=data_transform)\n",
    "train_dataset_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=2)\n",
    "print('train_dataset_loader',len(train_dataset_loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "oKyZ1aG7FMKv",
    "outputId": "b3ca247d-2284-42e4-fd98-3d292ccf2aeb"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test_dataset_loader 48\n"
     ]
    }
   ],
   "source": [
    "data_transform_test = transforms.Compose([#transforms.RandomResizedCrop(224),\n",
    "                                         #transforms.RandomHorizontalFlip(),\n",
    "                                         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",
    "test_dataset = datasets.ImageFolder(root='./DATA/output/test',transform = data_transform_test )\n",
    "#test_dataset = datasets.ImageFolder(root='/content/gdrive/My Drive/iitropar/ANN/P29-Cross-Pure-Dogs/test',transform = data_transform_test )\n",
    "test_dataset_loader = torch.utils.data.DataLoader(test_dataset, batch_size=32)\n",
    "print('test_dataset_loader',len(test_dataset_loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "PwyJBUh2FMKw"
   },
   "outputs": [],
   "source": [
    "#def sparse_categorical_accuracy(y_true, y_pred):\n",
    "#    return K.cast(K.equal(K.max(y_true, axis=-1), K.cast(K.argmax(y_pred, axis=-1), K.floatx())), K.floatx())\n",
    "\n",
    "def sparse_categorical_accuracy(y_true, y_pred):\n",
    "    #max_y_true,_= torch.max(y_true , dim = -1)\n",
    "    max_y_true = y_true\n",
    "    result = torch.eq(max_y_true , torch.argmax(y_pred, dim=-1).float())\n",
    "    result = result.float()\n",
    "    return result.sum().item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JNiJnBE9k-kW",
    "outputId": "d45a3886-1b4f-449c-d305-88b6b532f120"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ApplyMACE(\n",
       "  (basemodel): VGG(\n",
       "    (features): Sequential(\n",
       "      (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (1): ReLU(inplace=True)\n",
       "      (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (3): ReLU(inplace=True)\n",
       "      (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "      (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (6): ReLU(inplace=True)\n",
       "      (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (8): ReLU(inplace=True)\n",
       "      (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "      (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (11): ReLU(inplace=True)\n",
       "      (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (13): ReLU(inplace=True)\n",
       "      (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (15): ReLU(inplace=True)\n",
       "      (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "      (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (18): ReLU(inplace=True)\n",
       "      (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (20): ReLU(inplace=True)\n",
       "      (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (22): ReLU(inplace=True)\n",
       "      (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "      (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (25): ReLU(inplace=True)\n",
       "      (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (27): ReLU(inplace=True)\n",
       "      (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (29): ReLU(inplace=True)\n",
       "      (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    )\n",
       "    (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))\n",
       "    (classifier): Sequential(\n",
       "      (0): Linear(in_features=25088, out_features=4096, bias=True)\n",
       "      (1): ReLU(inplace=True)\n",
       "      (2): Dropout(p=0.5, inplace=False)\n",
       "      (3): Linear(in_features=4096, out_features=4096, bias=True)\n",
       "      (4): ReLU(inplace=True)\n",
       "      (5): Dropout(p=0.5, inplace=False)\n",
       "      (6): Linear(in_features=4096, out_features=20, bias=True)\n",
       "    )\n",
       "  )\n",
       "  (interpret_layer): MACENetwork(\n",
       "    (feature_extract): Sequential(\n",
       "      (0): Conv2d(512, 200, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (1): ReLU()\n",
       "    )\n",
       "    (protonet): ModuleList(\n",
       "      (0): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (1): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (2): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (3): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (4): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (5): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (6): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (7): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (8): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (9): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (10): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (11): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (12): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (13): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (14): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (15): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (16): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (17): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (18): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (19): EmbeddingGenerator(\n",
       "        (embeddings): Sequential(\n",
       "          (0): Linear(in_features=1, out_features=32, bias=True)\n",
       "          (1): Tanh()\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (concept_relevance): ModuleList(\n",
       "      (0): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (1): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (2): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (3): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (4): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (5): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (6): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (7): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (8): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (9): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (10): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (11): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (12): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (13): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (14): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (15): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (16): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (17): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (18): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "      (19): Relevance(\n",
       "        (rel): Linear(in_features=320, out_features=1, bias=True)\n",
       "      )\n",
       "    )\n",
       "    (softmax): Sigmoid()\n",
       "    (final_dense): Sequential(\n",
       "      (0): Linear(in_features=6400, out_features=4096, bias=True)\n",
       "      (1): ReLU()\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "meta_model = ApplyMACE()\n",
    "meta_model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0ORezYeyFMKw",
    "outputId": "220eecdd-6da0-4050-9169-a05408eebb3f"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "20"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "id": "3MTuznSKFMKw"
   },
   "outputs": [],
   "source": [
    "mse = torch.nn.MSELoss()\n",
    "kl = torch.nn.BCELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "id": "IIP0j1bZFMKw"
   },
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(meta_model.parameters(),lr=5e-4)\n",
    "#optimizer = optim.SGD(meta_model.parameters(), lr=1.0, momentum=0.9)\n",
    "#exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "id": "Bx5syZulFMKw"
   },
   "outputs": [],
   "source": [
    "train_acc = 0\n",
    "test_acc = 0\n",
    "rev_acc = 0\n",
    "train_total = 0\n",
    "test_total = 0\n",
    "\n",
    "train_mse_dense = 0\n",
    "train_kl = 0\n",
    "train_revelance = 0\n",
    "train_loss_triplet = 0\n",
    "num_batch=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "id": "WOAQRc8_arkj"
   },
   "outputs": [],
   "source": [
    "def train_step(images, labels):\n",
    "    global train_acc, train_total, train_mse_dense, train_kl, train_revelance, train_loss_triplet, num_batch  \n",
    "    \n",
    "    class_prototypes, class_protoweight, relevances, dense_layer_predict, first_dense, final_predict, final_output, class_protoweighted = meta_model(images)      \n",
    "    optimizer.zero_grad()\n",
    "    \n",
    "    # MSE with pre dense layer\n",
    "    mse_loss = mse(first_dense, dense_layer_predict)\n",
    "    \n",
    "    # KL with original distribution\n",
    "    kl_loss = kl(final_predict , final_output)\n",
    "    revelance_loss = kl( relevances , final_output)\n",
    "\n",
    "    triplet_loss = 0\n",
    "    for cp in class_prototypes:\n",
    "        # (batch_size, embedding_size, num_prototypes) (64,32,10)\n",
    "        prototype = torch.stack(cp, dim = 2)\n",
    "        triplet_loss += 0.1*new_triplet_loss(prototype.permute(0,2,1), 1)\n",
    "\n",
    "    triplet_loss = torch.mean(triplet_loss, dim=0)\n",
    "\n",
    "    # TOTAL LOSS\n",
    "    loss = mse_loss + kl_loss + revelance_loss + triplet_loss\n",
    "\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    #exp_lr_scheduler.step()\n",
    "    \n",
    "    train_mse_dense += mse_loss.item()\n",
    "    train_kl += kl_loss.item()\n",
    "    train_revelance += revelance_loss.item()\n",
    "    train_loss_triplet += triplet_loss.item()\n",
    "    num_batch += 1\n",
    "    \n",
    "    train_acc += sparse_categorical_accuracy(torch.argmax(final_output,dim=1), final_predict)\n",
    "    train_total += labels.size(0)\n",
    "    \n",
    "def test_step(images, labels):\n",
    "    global test_acc, rev_acc, test_total \n",
    "    \n",
    "    class_prototypes, class_protoweight, relevances, dense_layer_predict, first_dense, final_predict, final_output, class_protoweighted = meta_model(images)       \n",
    "    test_acc += sparse_categorical_accuracy(torch.argmax(final_output,dim=1), final_predict)\n",
    "    rev_acc += sparse_categorical_accuracy(torch.argmax(final_output,dim=1), relevances)\n",
    "    test_total += labels.size(0)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DEFWE-z6csu2",
    "outputId": "61c5a422-efea-4283-d3b1-e198e97405ff",
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Eh 1,MSE Loss:0.225, KL Loss:0.195, Triplet Loss:0.271, Relevance Loss:0.253, Train Acc:0.110, Test Acc:0.099, Rev Acc:0.058\n",
      "Eh 2,MSE Loss:0.222, KL Loss:0.197, Triplet Loss:0.125, Relevance Loss:0.198, Train Acc:0.085, Test Acc:0.068, Rev Acc:0.062\n",
      "Eh 3,MSE Loss:0.220, KL Loss:0.193, Triplet Loss:0.154, Relevance Loss:0.198, Train Acc:0.113, Test Acc:0.147, Rev Acc:0.040\n",
      "Eh 4,MSE Loss:0.223, KL Loss:0.160, Triplet Loss:0.075, Relevance Loss:0.200, Train Acc:0.271, Test Acc:0.467, Rev Acc:0.085\n",
      "Eh 5,MSE Loss:0.217, KL Loss:0.104, Triplet Loss:0.026, Relevance Loss:0.200, Train Acc:0.604, Test Acc:0.801, Rev Acc:0.104\n",
      "Eh 6,MSE Loss:0.208, KL Loss:0.060, Triplet Loss:0.059, Relevance Loss:0.189, Train Acc:0.813, Test Acc:0.896, Rev Acc:0.263\n",
      "Eh 7,MSE Loss:0.204, KL Loss:0.046, Triplet Loss:0.074, Relevance Loss:0.177, Train Acc:0.858, Test Acc:0.899, Rev Acc:0.390\n",
      "Eh 8,MSE Loss:0.193, KL Loss:0.041, Triplet Loss:0.124, Relevance Loss:0.163, Train Acc:0.871, Test Acc:0.915, Rev Acc:0.609\n",
      "Eh 9,MSE Loss:0.192, KL Loss:0.039, Triplet Loss:0.138, Relevance Loss:0.157, Train Acc:0.878, Test Acc:0.899, Rev Acc:0.528\n",
      "Eh 10,MSE Loss:0.189, KL Loss:0.039, Triplet Loss:0.103, Relevance Loss:0.160, Train Acc:0.877, Test Acc:0.910, Rev Acc:0.497\n",
      "Eh 11,MSE Loss:0.187, KL Loss:0.038, Triplet Loss:0.067, Relevance Loss:0.157, Train Acc:0.879, Test Acc:0.919, Rev Acc:0.542\n",
      "Eh 12,MSE Loss:0.184, KL Loss:0.034, Triplet Loss:0.089, Relevance Loss:0.150, Train Acc:0.891, Test Acc:0.920, Rev Acc:0.610\n",
      "Eh 13,MSE Loss:0.182, KL Loss:0.035, Triplet Loss:0.099, Relevance Loss:0.149, Train Acc:0.891, Test Acc:0.913, Rev Acc:0.601\n",
      "Eh 14,MSE Loss:0.179, KL Loss:0.035, Triplet Loss:0.080, Relevance Loss:0.148, Train Acc:0.893, Test Acc:0.925, Rev Acc:0.653\n",
      "Eh 15,MSE Loss:0.175, KL Loss:0.032, Triplet Loss:0.078, Relevance Loss:0.143, Train Acc:0.900, Test Acc:0.925, Rev Acc:0.677\n",
      "Eh 16,MSE Loss:0.172, KL Loss:0.032, Triplet Loss:0.091, Relevance Loss:0.137, Train Acc:0.901, Test Acc:0.927, Rev Acc:0.698\n",
      "Eh 17,MSE Loss:0.168, KL Loss:0.031, Triplet Loss:0.096, Relevance Loss:0.131, Train Acc:0.904, Test Acc:0.921, Rev Acc:0.751\n",
      "Eh 18,MSE Loss:0.161, KL Loss:0.030, Triplet Loss:0.084, Relevance Loss:0.125, Train Acc:0.909, Test Acc:0.931, Rev Acc:0.836\n",
      "Eh 19,MSE Loss:0.157, KL Loss:0.029, Triplet Loss:0.090, Relevance Loss:0.120, Train Acc:0.911, Test Acc:0.934, Rev Acc:0.849\n",
      "Eh 20,MSE Loss:0.152, KL Loss:0.029, Triplet Loss:0.089, Relevance Loss:0.114, Train Acc:0.913, Test Acc:0.937, Rev Acc:0.849\n",
      "Eh 21,MSE Loss:0.147, KL Loss:0.029, Triplet Loss:0.121, Relevance Loss:0.110, Train Acc:0.911, Test Acc:0.937, Rev Acc:0.872\n",
      "Eh 22,MSE Loss:0.145, KL Loss:0.030, Triplet Loss:0.116, Relevance Loss:0.109, Train Acc:0.907, Test Acc:0.929, Rev Acc:0.882\n",
      "Eh 23,MSE Loss:0.145, KL Loss:0.030, Triplet Loss:0.078, Relevance Loss:0.108, Train Acc:0.906, Test Acc:0.931, Rev Acc:0.888\n",
      "Eh 24,MSE Loss:0.163, KL Loss:0.043, Triplet Loss:0.098, Relevance Loss:0.109, Train Acc:0.877, Test Acc:0.774, Rev Acc:0.746\n",
      "Eh 25,MSE Loss:104.632, KL Loss:4.239, Triplet Loss:0.614, Relevance Loss:5.172, Train Acc:0.098, Test Acc:0.044, Rev Acc:0.056\n",
      "Eh 26,MSE Loss:27.599, KL Loss:4.660, Triplet Loss:0.520, Relevance Loss:7.104, Train Acc:0.050, Test Acc:0.060, Rev Acc:0.058\n",
      "Eh 27,MSE Loss:8.209, KL Loss:4.320, Triplet Loss:0.468, Relevance Loss:7.245, Train Acc:0.054, Test Acc:0.042, Rev Acc:0.058\n",
      "Eh 28,MSE Loss:3.255, KL Loss:2.428, Triplet Loss:0.453, Relevance Loss:7.259, Train Acc:0.057, Test Acc:0.052, Rev Acc:0.059\n",
      "Eh 29,MSE Loss:1.510, KL Loss:1.092, Triplet Loss:0.467, Relevance Loss:7.241, Train Acc:0.069, Test Acc:0.110, Rev Acc:0.059\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-27-76dbb26782f0>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     12\u001b[0m         \u001b[0mimage\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mimage\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m         \u001b[0mlabel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlabel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m         \u001b[0mtrain_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimage\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlabel\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     15\u001b[0m     \u001b[0;31m#print('train step done')\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-26-f287d761c84e>\u001b[0m in \u001b[0;36mtrain_step\u001b[0;34m(images, labels)\u001b[0m\n\u001b[1;32m     16\u001b[0m         \u001b[0;31m# (batch_size, embedding_size, num_prototypes) (64,32,10)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m         \u001b[0mprototype\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcp\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 18\u001b[0;31m         \u001b[0mtriplet_loss\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;36m0.1\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mnew_triplet_loss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprototype\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpermute\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     19\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     20\u001b[0m     \u001b[0mtriplet_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmean\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtriplet_loss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "train_kl_loss = []\n",
    "train_mse_dense_loss = []\n",
    "train_relevance_loss = []\n",
    "train_loss_triplet_loss = []\n",
    "\n",
    "best_model_wts = copy.deepcopy(meta_model.state_dict())\n",
    "best_acc = 0.0\n",
    "\n",
    "EPOCHS = 30\n",
    "for epoch in range(EPOCHS):\n",
    "    for image, label in train_dataset_loader:\n",
    "        image = image.to(device)\n",
    "        label = label.to(device)\n",
    "        train_step(image, label)        \n",
    "    #print('train step done')\n",
    "    \n",
    "    for image, label in test_dataset_loader:\n",
    "        image = image.to(device)\n",
    "        label = label.to(device)\n",
    "        test_step(image, label)    \n",
    "    #print('test step done')\n",
    "\n",
    "    train_kl_loss.append(train_kl/num_batch)\n",
    "    train_mse_dense_loss.append(train_mse_dense/num_batch)\n",
    "    train_relevance_loss.append(train_revelance/num_batch)\n",
    "    train_loss_triplet_loss.append(train_loss_triplet/num_batch)\n",
    "    \n",
    "    template = 'Eh {},MSE Loss:{:.3f}, KL Loss:{:.3f}, Triplet Loss:{:.3f}, Relevance Loss:{:.3f}, Train Acc:{:.3f}, Test Acc:{:.3f}, Rev Acc:{:.3f}'\n",
    "    print(template.format(epoch+1,\n",
    "                        train_mse_dense/num_batch,\n",
    "                        train_kl/num_batch,\n",
    "                        train_loss_triplet/num_batch,\n",
    "                        train_revelance/num_batch,\n",
    "                        train_acc/train_total,\n",
    "                        test_acc/test_total,\n",
    "                        rev_acc/test_total ))\n",
    "    \n",
    "    if (train_acc/train_total) > best_acc:\n",
    "        best_acc = train_acc/train_total\n",
    "        best_model_wts = copy.deepcopy(meta_model.state_dict())\n",
    "        \n",
    "   # Reset the metrics for the next epoch    \n",
    "    train_acc = 0\n",
    "    test_acc = 0\n",
    "    rev_acc = 0\n",
    "    train_total = 0\n",
    "    test_total = 0\n",
    "    \n",
    "    train_mse_dense = 0\n",
    "    train_kl = 0\n",
    "    train_revelance = 0\n",
    "    train_loss_triplet = 0\n",
    "    num_batch=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.9228643216080402\n"
     ]
    }
   ],
   "source": [
    "#lr=5e-3\n",
    "best = best_acc\n",
    "print(best)\n",
    "PATH = './max_pool_92acc.pth'\n",
    "torch.save({'model_state_dict':best_model_wts,'optimizer_state_dict':optimizer.state_dict()}, PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Xp0PFPoKFMKx"
   },
   "source": [
    "```\n",
    "concept_input torch.Size([64, 512, 14, 14])\n",
    "CM torch.Size([64, 200, 14, 14])\n",
    "CMlen 200\n",
    "CM_embed_input -> torch.Size([64, 14, 14])\n",
    "Embed_input -> torch.Size([64, 196])\n",
    "embed_out -> torch.Size([64, 32])\n",
    "prototype torch.Size([64, 32])\n",
    "cprototype shape -> 10 torch.Size([64, 32])\n",
    "cps shape-> torch.Size([64, 320])\n",
    "relev--> torch.Size([64, 1])\n",
    "cprototype shape -> 10 torch.Size([64, 32])\n",
    "R -> torch.Size([64, 20])\n",
    "expanded R  -> torch.Size([64, 20, 1])\n",
    "1 concept (before cat) torch.Size([64, 320])\n",
    "concepts after cat -> torch.Size([64, 6400])```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "y2EaT-pMKC1c"
   },
   "outputs": [],
   "source": []
  }
 ],
 "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
}