ResNetCaps_E.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function, division\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "from torch.optim import Adam\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import csv\n",
    "from skimage import io\n",
    "\n",
    "from PIL import Image\n",
    "import pandas as pd\n",
    "\n",
    "import numpy as np\n",
    "import torchvision\n",
    "from torchvision import datasets, models, transforms\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.autograd import Variable\n",
    "\n",
    "\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "import os\n",
    "import copy\n",
    "\n",
    "verbose = False\n",
    "USE_CUDA = True\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvLayer(nn.Module):\n",
    "    def __init__(self, in_channels=128,  out_channels=256, kernel_size=9):\n",
    "           \n",
    "        super(ConvLayer, self).__init__()\n",
    "\n",
    "        self.conv = nn.Conv2d(in_channels=in_channels,\n",
    "                               out_channels=out_channels,\n",
    "                               kernel_size=kernel_size,\n",
    "                               stride=1\n",
    "                             )\n",
    "\n",
    "    def forward(self, x):\n",
    "        if verbose: print( \"Conv input size{}\".format(x.size()))\n",
    "        output = F.relu(self.conv(x))\n",
    "        if verbose: print(\"Conv output feature matrix {}\".format(output.shape))\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PrimaryCaps(nn.Module):\n",
    "    def __init__(self, num_capsules=8, in_channels=256, out_channels=32, kernel_size=9):\n",
    "\n",
    "        super(PrimaryCaps, self).__init__()\n",
    "\n",
    "        self.capsules = nn.ModuleList([\n",
    "            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=2, padding=0) \n",
    "                          for _ in range(num_capsules)])\n",
    "    \n",
    "    def forward(self, x, dimension = 32*6*6):\n",
    "        u = [capsule(x) for capsule in self.capsules]\n",
    "        u = torch.stack(u, dim=1)\n",
    "        if verbose: print( \"PrimaryCaps {}\".format(u.size()))\n",
    "        u = u.view(x.size(0), dimension, -1)\n",
    "        if verbose: print(\"PrimaryCaps size U {}\".format(u.size()))\n",
    "        output = self.squash(u)\n",
    "        if verbose: print(\"Primary Caps output {}\".format(output.size()))\n",
    "        return output\n",
    "    \n",
    "    def squash(self, input_tensor):\n",
    "        squared_norm = (input_tensor ** 2).sum(-1, keepdim=True)\n",
    "        output_tensor = squared_norm *  input_tensor / ((1. + squared_norm) * torch.sqrt(squared_norm))\n",
    "        if verbose: print(output_tensor.size())\n",
    "        return output_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DigitCaps(nn.Module):\n",
    "    def __init__(self, num_capsules = 1, num_routes=32 * 6 * 6 , in_channels=8,  out_channels=16):\n",
    "        \n",
    "        super(DigitCaps, self).__init__()\n",
    "\n",
    "        self.in_channels = in_channels\n",
    "        self.num_routes = num_routes\n",
    "        self.num_capsules = num_capsules\n",
    "\n",
    "        self.W = nn.Parameter(torch.randn(1, num_routes, num_capsules, out_channels, in_channels))\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.size(0)\n",
    "        x = torch.stack([x] * self.num_capsules, dim=2).unsqueeze(4)\n",
    "        if verbose: print( \"DigitCaps x {}, W {}\".format(x.size(),self.W.size()))\n",
    "        W = torch.cat([self.W] * batch_size, dim=0)\n",
    "        if verbose: print(\"DigitCaps W {}\".format(W.size()))\n",
    "        u_hat = torch.matmul(W, x)\n",
    "        if verbose: print(\"DigitCaps u_hat {}\".format(u_hat.size()))\n",
    "        \n",
    "        b_ij = Variable(torch.zeros(1, self.num_routes, self.num_capsules, 1))\n",
    "        b_ij = b_ij.to(device)#cuda()\n",
    "\n",
    "        num_iterations = 3\n",
    "        for iteration in range(num_iterations):\n",
    "            c_ij = F.softmax(b_ij)\n",
    "            c_ij = torch.cat([c_ij] * batch_size, dim=0).unsqueeze(4)\n",
    "\n",
    "            s_j = (c_ij * u_hat).sum(dim=1, keepdim=True)\n",
    "            v_j = self.squash(s_j)\n",
    "            \n",
    "            if iteration < num_iterations - 1:\n",
    "                a_ij = torch.matmul(u_hat.transpose(3, 4), torch.cat([v_j] * self.num_routes, dim=1))\n",
    "                b_ij = b_ij + a_ij.squeeze(4).mean(dim=0, keepdim=True)\n",
    "\n",
    "        return v_j.squeeze(1)\n",
    "    \n",
    "    def squash(self, input_tensor):\n",
    "        squared_norm = (input_tensor ** 2).sum(-1, keepdim=True)\n",
    "        output_tensor = squared_norm *  input_tensor / ((1. + squared_norm) * torch.sqrt(squared_norm))\n",
    "        return output_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResNetCaps_E(nn.Module):\n",
    "    def __init__(self,DigitEnd=True):\n",
    "        super(ResNetCaps_E, self).__init__()\n",
    "        self.DigitEnd = DigitEnd\n",
    "\n",
    "        model = torchvision.models.resnet18(pretrained=True)\n",
    "        modules = list(model.children())[:-4]\n",
    "        self.model=nn.Sequential(*modules)\n",
    "        for param in self.model.parameters():\n",
    "            param.requires_grad = False\n",
    "        if DigitEnd:\n",
    "            self.model.layer3 = nn.Sequential(ConvLayer(), PrimaryCaps(),DigitCaps())\n",
    "        else:\n",
    "            self.model.layer3 = nn.Sequential(ConvLayer(), PrimaryCaps())    \n",
    "                                          \n",
    "    def forward(self,x):\n",
    "        \n",
    "        output = self.model(x)\n",
    "        if not self.DigitEnd:\n",
    "            output = maxpooling_(output, modal='capsules')\n",
    "        return output\n",
    "\n",
    "    def maxpooling_(x, modal='capsules'):\n",
    "        m = []\n",
    "        if modal == 'capsules':\n",
    "            k = torch.argmax(x,dim=2)\n",
    "            for i in range(x.size(0)):\n",
    "                for j in range(x.size(1)):\n",
    "                    m.append(x[i,j,k[i,j].item()])\n",
    "            m = (torch.Tensor(m).view(x.size(0),x.size(1)))\n",
    "        else:\n",
    "            pooling_layer = nn.MaxPool2d((3,8),stride=2)\n",
    "            m = pooling_layer(x)\n",
    "        m.requires_grad = True            \n",
    "        return m\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}