CNN_2DselfAttention.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision import models\n",
    "from torchsummary import summary \n",
    "from torchvision import models\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading CIFAR10 dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "batch_size = 64\n",
    "all_transforms = transforms.Compose([transforms.Resize((32,32)),\n",
    "                                     transforms.ToTensor(),\n",
    "                                     transforms.Normalize(mean=[0.4914, 0.4822, 0.4465],\n",
    "                                                          std=[0.2023, 0.1994, 0.2010])\n",
    "                                     ])\n",
    "# Create Training dataset\n",
    "train_dataset = torchvision.datasets.CIFAR10(root = './data',\n",
    "                                             train = True,\n",
    "                                             transform = all_transforms,\n",
    "                                             download = True)\n",
    "\n",
    "# Create Testing dataset\n",
    "test_dataset = torchvision.datasets.CIFAR10(root = './data',\n",
    "                                            train = False,\n",
    "                                            transform = all_transforms,\n",
    "                                            download=True)\n",
    "\n",
    "# Instantiate loader objects to facilitate processing\n",
    "train_loader = torch.utils.data.DataLoader(dataset = train_dataset,\n",
    "                                           batch_size = batch_size,\n",
    "                                           shuffle = True)\n",
    "\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(dataset = test_dataset,\n",
    "                                           batch_size = batch_size,\n",
    "                                           shuffle = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2D SelfAttention model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Source of attention model is \n",
    "__all__ = [\"SelfAttention2d\"]\n",
    "\n",
    "\n",
    "class SelfAttention2d(nn.Module):\n",
    "    r\"\"\"Self Attention Module as proposed in the paper `\"Self-Attention Generative Adversarial\n",
    "    Networks by Han Zhang et. al.\" <https://arxiv.org/abs/1805.08318>`_\n",
    "    .. math:: attention = softmax((query(x))^T * key(x))\n",
    "    .. math:: output = \\gamma * value(x) * attention + x\n",
    "    where\n",
    "    - :math:`query` : 2D Convolution Operation\n",
    "    - :math:`key` : 2D Convolution Operation\n",
    "    - :math:`value` : 2D Convolution Operation\n",
    "    - :math:`x` : Input\n",
    "    Args:\n",
    "        input_dims (int): The input channel dimension in the input ``x``.\n",
    "        output_dims (int, optional): The output channel dimension. If ``None`` the output\n",
    "            channel value is computed as ``input_dims // 8``. So if the ``input_dims`` is **less\n",
    "            than 8** then the layer will give an error.\n",
    "        return_attn (bool, optional): Set it to ``True`` if you want the attention values to be\n",
    "            returned.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, input_dims, output_dims=None, return_attn=False):\n",
    "        output_dims = input_dims // 8 if output_dims is None else output_dims\n",
    "        if output_dims == 0:\n",
    "            raise Exception(\n",
    "                \"The output dims corresponding to the input dims is 0. Increase the input\\\n",
    "                            dims to 8 or more. Else specify output_dims\"\n",
    "            )\n",
    "        super(SelfAttention2d, self).__init__()\n",
    "        self.query = nn.Conv2d(input_dims, output_dims, 1)\n",
    "        self.key = nn.Conv2d(input_dims, output_dims, 1)\n",
    "        self.value = nn.Conv2d(input_dims, input_dims, 1)\n",
    "        self.gamma = nn.Parameter(torch.zeros(1))\n",
    "        self.return_attn = return_attn\n",
    "\n",
    "    def forward(self, x):\n",
    "        #print(x.shape)  # [bs,64,5,5]\n",
    "        r\"\"\"Computes the output of the Self Attention Layer\n",
    "        Args:\n",
    "            x (torch.Tensor): A 4D Tensor with the channel dimension same as ``input_dims``.\n",
    "        Returns:\n",
    "            A tuple of the ``output`` and the ``attention`` if ``return_attn`` is set to ``True``\n",
    "            else just the ``output`` tensor.\n",
    "        \"\"\"\n",
    "        dims = (x.size(0), -1, x.size(2) * x.size(3)) # [2,-1,5*5]\n",
    "        out_query = self.query(x).view(dims)\n",
    "        #print('this is the query',out_query.shape) #  torch.Size([2, 8, 25])\n",
    "        out_key = self.key(x).view(dims).permute(0, 2, 1) \n",
    "        #print('this is the key',out_key.shape) #  torch.Size([2, 25, 8])\n",
    "        attn = F.softmax(torch.bmm(out_key, out_query), dim=-1) #  torch.Size([2, 25, 25])\n",
    "        #print('shape of attn var', attn.shape)\n",
    "        out_value = self.value(x).view(dims)\n",
    "        #print('shape of out_value var', out_value.shape)\n",
    "\n",
    "        out_value = torch.bmm(out_value, attn).view(x.size())\n",
    "        out = self.gamma * out_value + x\n",
    "        if self.return_attn:\n",
    "            return out, attn\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1           [-1, 32, 30, 30]             896\n",
      "            Conv2d-2           [-1, 32, 28, 28]           9,248\n",
      "         MaxPool2d-3           [-1, 32, 14, 14]               0\n",
      "            Conv2d-4           [-1, 64, 12, 12]          18,496\n",
      "            Conv2d-5           [-1, 64, 10, 10]          36,928\n",
      "         MaxPool2d-6             [-1, 64, 5, 5]               0\n",
      "            Conv2d-7              [-1, 8, 5, 5]             520\n",
      "            Conv2d-8              [-1, 8, 5, 5]             520\n",
      "            Conv2d-9             [-1, 64, 5, 5]           4,160\n",
      "  SelfAttention2d-10             [-1, 64, 5, 5]               0\n",
      "      BatchNorm2d-11             [-1, 64, 5, 5]             128\n",
      "           Linear-12                  [-1, 128]         204,928\n",
      "          Dropout-13                  [-1, 128]               0\n",
      "             ReLU-14                  [-1, 128]               0\n",
      "           Linear-15                   [-1, 10]           1,290\n",
      "================================================================\n",
      "Total params: 277,114\n",
      "Trainable params: 277,114\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.01\n",
      "Forward/backward pass size (MB): 0.63\n",
      "Params size (MB): 1.06\n",
      "Estimated Total Size (MB): 1.70\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "\n",
    "\n",
    "# Creating a CNN class\n",
    "class CNN_At(nn.Module):\n",
    "\t#  Determine what layers and their order in CNN object \n",
    "    def __init__(self, num_classes=10):\n",
    "        super(CNN_At, self).__init__()\n",
    "        self.conv_layer1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3,padding = 0) \n",
    "        self.conv_layer2 = nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3)\n",
    "        self.max_pool1 = nn.MaxPool2d(kernel_size = 2, stride = 2)\n",
    "        \n",
    "        self.conv_layer3 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)\n",
    "        self.conv_layer4 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3)\n",
    "        self.max_pool2 = nn.MaxPool2d(kernel_size = 2, stride = 2)\n",
    "        self.atten = SelfAttention2d(64)\n",
    "        self.bn = nn.BatchNorm2d(64)\n",
    "\n",
    "        \n",
    "        self.fc1 = nn.Linear(1600, 128)# [64 * 5 * 5]\n",
    "        self.dropout = nn.Dropout(p=0.1)\n",
    "\n",
    "       \n",
    "        self.relu1 = nn.ReLU()\n",
    "        \n",
    "        self.fc2 = nn.Linear(128, num_classes)\n",
    "    \n",
    "    # Progresses data across layers    \n",
    "    def forward(self, x):\n",
    "        out = self.conv_layer1(x)\n",
    "        out = self.conv_layer2(out)\n",
    "        out = self.max_pool1(out)\n",
    "        \n",
    "        out = self.conv_layer3(out)\n",
    "        out = self.conv_layer4(out)\n",
    "        out = self.max_pool2(out)\n",
    "        out = self.atten(out)\n",
    "        out = self.bn(out)\n",
    "\n",
    "                \n",
    "        out = out.reshape(out.size(0), -1)\n",
    "        \n",
    "        out = self.fc1(out)\n",
    "        out = self.dropout(out)\n",
    "\n",
    "        out = self.relu1(out)\n",
    "        \n",
    "        out = self.fc2(out)\n",
    "        return out\n",
    "    \n",
    "    \n",
    "model_atten = CNN_At()\n",
    "model_atten = model_atten.to(device=device,dtype=torch.float)\n",
    "summary(model_atten,(3,32,32))\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/20], Loss: 1.0352\n",
      "Epoch [2/20], Loss: 1.1540\n",
      "Epoch [3/20], Loss: 0.6718\n",
      "Epoch [4/20], Loss: 0.8419\n",
      "Epoch [5/20], Loss: 0.5113\n",
      "Epoch [6/20], Loss: 1.0275\n",
      "Epoch [7/20], Loss: 1.5056\n",
      "Epoch [8/20], Loss: 0.6464\n",
      "Epoch [9/20], Loss: 0.5607\n",
      "Epoch [10/20], Loss: 0.4222\n",
      "Epoch [11/20], Loss: 0.3412\n",
      "Epoch [12/20], Loss: 0.7515\n",
      "Epoch [13/20], Loss: 0.4577\n",
      "Epoch [14/20], Loss: 0.7822\n",
      "Epoch [15/20], Loss: 0.8287\n",
      "Epoch [16/20], Loss: 0.2697\n",
      "Epoch [17/20], Loss: 0.3677\n",
      "Epoch [18/20], Loss: 0.3275\n",
      "Epoch [19/20], Loss: 0.4244\n",
      "Epoch [20/20], Loss: 0.4315\n"
     ]
    }
   ],
   "source": [
    "batch_size = 64\n",
    "num_classes = 10\n",
    "learning_rate = 0.001\n",
    "num_epochs = 20\n",
    "\n",
    "model_atten = model_atten.to(device=device,dtype=torch.float)\n",
    "# Set Loss function with criterion\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "# Set optimizer with optimizer\n",
    "#optimizer = torch.optim.SGD(model_atten.parameters(), lr=learning_rate, weight_decay = 0.005, momentum = 0.9)  \n",
    "#optimizer = torch.optim.Adam(model_atten.parameters(), lr=learning_rate, weight_decay = 0.005)  \n",
    "optimizer = torch.optim.Adam(model_atten.parameters(),lr=0.0001)\n",
    "total_step = len(train_loader)\n",
    "\n",
    "# We use the pre-defined number of epochs to determine how many iterations to train the network on\n",
    "for epoch in range(num_epochs):\n",
    "\t#Load in the data in batches using the train_loader object\n",
    "    for i, (images, labels) in enumerate(train_loader):  \n",
    "        # Move tensors to the configured device\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "        \n",
    "        # Forward pass\n",
    "        outputs = model_atten(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        \n",
    "        # Backward and optimize\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))\n",
    "\n",
    "       \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy of the network on the 50000 train images: 90.266 %\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    for images, labels in test_loader:\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "        outputs = model_atten(images)\n",
    "        _, predicted = torch.max(outputs.data, 1)\n",
    "        total += labels.size(0)\n",
    "        correct += (predicted == labels).sum().item()\n",
    "    \n",
    "    print('Accuracy of the network on the {} train images: {} %'.format(50000, 100 * correct / total))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "clamEnv",
   "language": "python",
   "name": "clamenv"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}