MTAN Implementation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "front-annual",
   "metadata": {},
   "source": [
    "# MTAN\n",
    "- MTAN is a state-of-the art architecture for Multi-task learning\n",
    "- The official implementation is a bit messy, here I will propose a simpler version and test that it produces the same outputs\n",
    "- I will provide the implementation for MTAN on VGG16 and Resnet backbones\n",
    "- To keep it simple enough for a notebook, these implementation will be for multi-domain classification rather than dense predictions\n",
    "- Extending it to dense predictions can be done by attaching decoders/deeplab heads and potentially keep track of pooling indices\n",
    "- Manage to significantly reduce size of VGG/ResNet code but using a dedicated attention module\n",
    "- This notebook is structured as follows: notes to MTAN, My simplified implementation, Tests against original copy-pasted source code."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "prescription-factory",
   "metadata": {},
   "source": [
    "![](https://vitalab.github.io/article/images/MTAN/architecture.png)\n",
    "\n",
    "Notes on MTAN:\n",
    "- Soft-parameter sharing method, so each task has it's own parameters/activations at every level -> and a different computational graph\n",
    "- There is one attention module per block/ stage\n",
    "- The attention module learns to putput a mask per task, which is then applied to the representations of the shared backbone\n",
    "- Each attention module has 3 inputs: two inputs from the shared backbone, and the output from the previous attention module\n",
    "- Each attention module itself has task-specific and shared parameters\n",
    "- The task specific attention parameters only compute the mask\n",
    "- The shared attention parameters act as an alternative backbone; hence they will be built similar to the backbone architecture\n",
    "- Inputs for the attention module are the intermediary outputs from each backbone block/stage. Where these outputs are taken from seems a bit arbitrary/architecture dependent. One input (on which the task specific mask is applied) is alwasy the output of the backbone stage. However the other input, the one that goes into the computation of the mask, can be takern either from the beginning or somewhere in the middle of the block/stage.\n",
    "- The task specific mask is computed on information from both the backbone and the previous attention -> these are concatenated hence the 2x input size of attention convs.\n",
    "- The reason the implmentation becomes messy is because of the inputs to the attention module: it requires outputs of conv layers / blocks from inside backbone stages. This is not great since one cannot just encapsulate the backbone, one has to run each layer manually and retreive the selected intermediary representation. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "brilliant-halifax",
   "metadata": {},
   "source": [
    "# Simple Implementation\n",
    "- The key is to define a dedicated attention module in a smart way\n",
    "- Becuase the shared part of the attention module is architecture dependent, it has to be an argument\n",
    "- Also generally spearking VGG and resnet architectures seem to be parametrised (as in having arch hyper-params) however the truth is that they have a lot of idiosyncracies which make them a lot less generalisable. Therefore sometimes the most elegat way is to just hardcode indices or stages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "mighty-memory",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import itertools\n",
    "from torch import nn\n",
    "from typing import List, Dict\n",
    "import torch.nn.functional as F\n",
    "from torchvision.models import vgg16_bn, resnet18, resnet50, resnet101\n",
    "from torchvision.models.resnet import BasicBlock, Bottleneck, conv1x1\n",
    "\n",
    "\n",
    "class ConvBNReLU(nn.Sequential):\n",
    "    \"\"\"Shorthand for conv layer\"\"\"\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super().__init__(\n",
    "            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(out_channels),\n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "\n",
    "def init_model(model: nn.Module):\n",
    "    for m in model.modules():\n",
    "        if isinstance(m, nn.Conv2d):\n",
    "            nn.init.xavier_normal_(m.weight)\n",
    "            if m.bias is not None:\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "        elif isinstance(m, nn.BatchNorm2d):\n",
    "            nn.init.constant_(m.weight, 1)\n",
    "            nn.init.constant_(m.bias, 0)\n",
    "        elif isinstance(m, nn.Linear):\n",
    "            nn.init.xavier_normal_(m.weight)\n",
    "            nn.init.constant_(m.bias, 0)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "posted-cartridge",
   "metadata": {},
   "source": [
    "## MTAN Attention\n",
    "- shared_attention sub-module has to be defined outside as it depends on the backbone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "blessed-disclaimer",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MTANAttention(nn.Module):\n",
    "    \"\"\"\n",
    "    Combines task specific attentions and shared parameters\n",
    "    Task specific params, which compute the attion mask, are always defined the same way\n",
    "    Shared parameters depend on architecture as they try to emulate compoents in the backbone\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, tasks: List[str], in_channels: int, mid_channels: int, out_channels: int, shared_att: nn.Module = None):\n",
    "        \"\"\"\n",
    "        :param in_channels: If this modules receives inputs from a previous attention module, then in_channels will have to be 2x the equivalent backbone input\n",
    "        :param shared_att: Architecture dependent - For VGG it's a 3x3 conv, for ResNets it's a bottleneck layer; Optional, if missing module outputs masked in2\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.task_att = nn.ModuleDict({\n",
    "            task: nn.Sequential(\n",
    "                    nn.Conv2d(in_channels=in_channels, out_channels=mid_channels, kernel_size=1, padding=0),\n",
    "                    nn.BatchNorm2d(mid_channels),\n",
    "                    nn.ReLU(inplace=True),\n",
    "                    nn.Conv2d(in_channels=mid_channels, out_channels=out_channels, kernel_size=1, padding=0),\n",
    "                    nn.BatchNorm2d(out_channels),\n",
    "                    nn.Sigmoid(),\n",
    "                )\n",
    "            for task in tasks\n",
    "        }) \n",
    "        self.shared_att = shared_att if shared_att is not None else nn.Sequential()\n",
    "            \n",
    "    def forward(self, task: str, in1: torch.Tensor, in2: torch.Tensor, att_in: torch.Tensor = None):\n",
    "        \"\"\"\n",
    "        :param in1: the output from somewhere in the shared backbone block/stage (first vertical line in figure)\n",
    "        :param in2: the output from the shared backbone  block/stage (second vertical line in figure)\n",
    "        :param att_in: Optional; the output from a previous attention module; \n",
    "        \"\"\"\n",
    "        out = in1\n",
    "        if att_in is not None:\n",
    "            out = torch.cat([out, att_in], dim=1) # merge output from previous attention module if available\n",
    "        out = self.task_att[task](out)  # The two task-specific convs\n",
    "        out = out * in2 # element-wise multiplication with second input (from shared blockss)\n",
    "        out = self.shared_att(out) # The shared conv at the end\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "compact-belarus",
   "metadata": {},
   "source": [
    "## VGG16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "saved-advertising",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MTANVGG16(nn.Module):\n",
    "    \"\"\"\n",
    "    For multi-domain classification. Attaches a classifier on top of outputs from last attention module.\n",
    "    Using vgg16_bn model as backbone from torchvision\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, tasks: Dict[str, int]):\n",
    "        \"\"\"\n",
    "        :param tasks: {task_name: n_classes} e.g {\"T1\": 5, \"T2\": 3}\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.tasks = tasks\n",
    "        self.shared_layers = vgg16_bn().features[:-1] # dropping the last maxpool\n",
    "        # Not worth trying to use loops\n",
    "        self.attentions = nn.ModuleList([\n",
    "            MTANAttention(tasks, 64, 64, 64, ConvBNReLU(64, 128)),\n",
    "            MTANAttention(tasks, 2 * 128, 128, 128, ConvBNReLU(128, 256)), # in_channels is twice as big due to concatenation\n",
    "            MTANAttention(tasks, 2 * 256, 256, 256, ConvBNReLU(256, 512)), \n",
    "            MTANAttention(tasks, 2 * 512, 512, 512, ConvBNReLU(512, 512)),\n",
    "            MTANAttention(tasks, 2 * 512, 512, 512, ConvBNReLU(512, 512)),\n",
    "        ])\n",
    "        self.classifier = nn.ModuleDict({task: nn.Linear(512, task_classes) for task, task_classes in self.tasks.items()})\n",
    "        \n",
    "    def forward(self, task: str, X):\n",
    "        \"\"\"\n",
    "        Classification output for a single task\n",
    "        \"\"\"\n",
    "        \n",
    "        # Need to keep activations from conv-bn-relu layers, not more granular than that\n",
    "        sh_outs = []\n",
    "        for layer in self.shared_layers:\n",
    "            X = layer(X)\n",
    "            if isinstance(X, nn.ReLU):\n",
    "                sh_outs.append(X)\n",
    "        \n",
    "        maxpool = lambda x: F.max_pool2d(x, kernel_size=2, stride=2)\n",
    "        sh_out_idx = [(0, 1), (2, 3), (4, 6), (7, 9), (10, 12)] # indices of the relevant outputs from the shared backbone that go into the attention module\n",
    "        att_out = None\n",
    "        for block, (in1_idx, in2_idx) in enumerate(sh_out_idx):\n",
    "            att_out = maxpool(self.attentions[block](task, sh_outs[in1_idx], sh_outs[in2_idx], att_out))\n",
    "            \n",
    "        out = torch.flatten(att_out, 1)\n",
    "        out = self.classifier[task](out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "equipped-royal",
   "metadata": {},
   "source": [
    "## ResNets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "processed-joseph",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MTANResNet(nn.Module):\n",
    "    \"\"\"\n",
    "    Should work for both Basic and Bottleneck based ResNets\n",
    "    Assuming the first input to the attention module is always the output of the penultimate block in a stage\n",
    "    Task-specific layers have a 4x bottleneck\n",
    "    Uses maxpool between stages\n",
    "    \"\"\"\n",
    "    def __init__(self, tasks: Dict[str, int], backbone=resnet50()):\n",
    "        \"\"\"\n",
    "        :param tasks: {task_name: n_classes} e.g {\"T1\": 5, \"T2\": 3}\n",
    "        :param backbone: yes, I know, default value not good practice. \n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.tasks = tasks\n",
    "        self.shared_conv = nn.Sequential(backbone.conv1, backbone.bn1, backbone.relu, backbone.maxpool)\n",
    "        \n",
    "        # splitting each stage in two for the in1 and in2 inputs; and storing as one list\n",
    "        self.shared_layers = nn.ModuleList(\n",
    "            itertools.chain.from_iterable(\n",
    "                [(layer[:-1], layer[-1]) \n",
    "                 for layer in [backbone.__getattr__(f\"layer{i}\") for i in range(1,5)]\n",
    "                ]\n",
    "        ))\n",
    "        \n",
    "        BlockCls = self.shared_layers[-1].__class__ # hopefuly BasicBlock or Bottleneck\n",
    "        shared_att = lambda ch_in, ch_out: BlockCls(ch_in, ch_out, downsample=nn.Sequential(conv1x1(ch_in, 4 * ch_out), nn.BatchNorm2d(4 * ch_out)))  \n",
    "        \n",
    "        ch = [l.conv1.in_channels for l in self.shared_layers[1::2]] # first conv of last block\n",
    "        self.attentions = nn.ModuleList([\n",
    "            MTANAttention(tasks, ch[0], ch[0] // 4, ch[0], shared_att(ch[0], ch[1] // 4)),\n",
    "            MTANAttention(tasks, 2 * ch[1], ch[1] // 4, ch[1], shared_att(ch[1], ch[2] // 4)),\n",
    "            MTANAttention(tasks, 2 * ch[2], ch[2] // 4, ch[2], shared_att(ch[2], ch[3] // 4)),\n",
    "            MTANAttention(tasks, 2 * ch[3], ch[3] // 4, ch[3]), # final attention doesnt have a shared part\n",
    "        ])\n",
    "        self.classifier = nn.ModuleDict({task: nn.Linear(ch[3], task_classes) for task, task_classes in self.tasks.items()})\n",
    "        \n",
    "    def forward(self, task: str, X):\n",
    "        \"\"\"\n",
    "        Classification output for a single task\n",
    "        \"\"\"\n",
    "        out = self.shared_conv(X)\n",
    "        sh_outs = [out := l(out) for l in self.shared_layers]\n",
    "        \n",
    "        maxpool = lambda x: F.max_pool2d(x, kernel_size=2, stride=2)\n",
    "        att_out = None\n",
    "        for stage in range(4):\n",
    "            att_out = self.attentions[stage](task, sh_outs[stage * 2], sh_outs[stage * 2 + 1], att_out)\n",
    "            if stage < 3:\n",
    "                # NOTE: This can/should be changed if backbone is different (e.g uses dilation)\n",
    "                att_out = maxpool(att_out)\n",
    "        \n",
    "        out = F.avg_pool2d(att_out, 8)\n",
    "        out = torch.flatten(out, 1)\n",
    "        out = self.classifier[task](out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "nominated-glory",
   "metadata": {},
   "source": [
    "# Tests\n",
    "- Instead of running benchmarks I will just compare the output of my implementation to the output of the original\n",
    "- Official MTAN implementation does not have a library structure, meaning I cannot import individual modules easily. I decided to copy-paste implementations and just comment out stuff that is not relevant to the comparison.\n",
    "- Copying layer weights rather than fiddling with the seed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "partial-motor",
   "metadata": {},
   "source": [
    "## VGG16\n",
    "- VGG16 for domain classification is actually never used in the paper, however we can compare against the Encoder part of SegNet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "quality-international",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Original Implementation\n",
    "- commenting out the decoder part\n",
    "- outputting only the last attentions from the encoder (which would feed into classifiers if there were any)\n",
    "- No other changes to the orignal code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "everyday-thomas",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class SegNet(nn.Module):\n",
    "    \"\"\"\n",
    "    This is the original implementaiton @https://github.com/lorenmt/mtan/blob/master/im2im_pred/model_segnet_mtan.py\n",
    "    \"\"\"\n",
    "    def __init__(self):\n",
    "        super(SegNet, self).__init__()\n",
    "        # initialise network parameters\n",
    "        filter = [64, 128, 256, 512, 512]\n",
    "        self.class_nb = 13\n",
    "\n",
    "        # define encoder decoder layers\n",
    "        self.encoder_block = nn.ModuleList([self.conv_layer([3, filter[0]])])\n",
    "        self.decoder_block = nn.ModuleList([self.conv_layer([filter[0], filter[0]])])\n",
    "        for i in range(4):\n",
    "            self.encoder_block.append(self.conv_layer([filter[i], filter[i + 1]]))\n",
    "            self.decoder_block.append(self.conv_layer([filter[i + 1], filter[i]]))\n",
    "\n",
    "        # define convolution layer\n",
    "        self.conv_block_enc = nn.ModuleList([self.conv_layer([filter[0], filter[0]])])\n",
    "        self.conv_block_dec = nn.ModuleList([self.conv_layer([filter[0], filter[0]])])\n",
    "        for i in range(4):\n",
    "            if i == 0:\n",
    "                self.conv_block_enc.append(self.conv_layer([filter[i + 1], filter[i + 1]]))\n",
    "                self.conv_block_dec.append(self.conv_layer([filter[i], filter[i]]))\n",
    "            else:\n",
    "                self.conv_block_enc.append(nn.Sequential(self.conv_layer([filter[i + 1], filter[i + 1]]),\n",
    "                                                         self.conv_layer([filter[i + 1], filter[i + 1]])))\n",
    "                self.conv_block_dec.append(nn.Sequential(self.conv_layer([filter[i], filter[i]]),\n",
    "                                                         self.conv_layer([filter[i], filter[i]])))\n",
    "\n",
    "        # define task attention layers\n",
    "        self.encoder_att = nn.ModuleList([nn.ModuleList([self.att_layer([filter[0], filter[0], filter[0]])])])\n",
    "        self.decoder_att = nn.ModuleList([nn.ModuleList([self.att_layer([2 * filter[0], filter[0], filter[0]])])])\n",
    "        self.encoder_block_att = nn.ModuleList([self.conv_layer([filter[0], filter[1]])])\n",
    "        self.decoder_block_att = nn.ModuleList([self.conv_layer([filter[0], filter[0]])])\n",
    "\n",
    "        for j in range(3):\n",
    "            if j < 2:\n",
    "                self.encoder_att.append(nn.ModuleList([self.att_layer([filter[0], filter[0], filter[0]])]))\n",
    "                self.decoder_att.append(nn.ModuleList([self.att_layer([2 * filter[0], filter[0], filter[0]])]))\n",
    "            for i in range(4):\n",
    "                self.encoder_att[j].append(self.att_layer([2 * filter[i + 1], filter[i + 1], filter[i + 1]]))\n",
    "                self.decoder_att[j].append(self.att_layer([filter[i + 1] + filter[i], filter[i], filter[i]]))\n",
    "\n",
    "        for i in range(4):\n",
    "            if i < 3:\n",
    "                self.encoder_block_att.append(self.conv_layer([filter[i + 1], filter[i + 2]]))\n",
    "                self.decoder_block_att.append(self.conv_layer([filter[i + 1], filter[i]]))\n",
    "            else:\n",
    "                self.encoder_block_att.append(self.conv_layer([filter[i + 1], filter[i + 1]]))\n",
    "                self.decoder_block_att.append(self.conv_layer([filter[i + 1], filter[i + 1]]))\n",
    "\n",
    "        self.pred_task1 = self.conv_layer([filter[0], self.class_nb], pred=True)\n",
    "        self.pred_task2 = self.conv_layer([filter[0], 1], pred=True)\n",
    "        self.pred_task3 = self.conv_layer([filter[0], 3], pred=True)\n",
    "\n",
    "        # define pooling and unpooling functions\n",
    "        self.down_sampling = nn.MaxPool2d(kernel_size=2, stride=2, return_indices=True)\n",
    "        self.up_sampling = nn.MaxUnpool2d(kernel_size=2, stride=2)\n",
    "\n",
    "        self.logsigma = nn.Parameter(torch.FloatTensor([-0.5, -0.5, -0.5]))\n",
    "\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Conv2d):\n",
    "                nn.init.xavier_normal_(m.weight)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "            elif isinstance(m, nn.Linear):\n",
    "                nn.init.xavier_normal_(m.weight)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "\n",
    "    def conv_layer(self, channel, pred=False):\n",
    "        if not pred:\n",
    "            conv_block = nn.Sequential(\n",
    "                nn.Conv2d(in_channels=channel[0], out_channels=channel[1], kernel_size=3, padding=1),\n",
    "                nn.BatchNorm2d(num_features=channel[1]),\n",
    "                nn.ReLU(inplace=True),\n",
    "            )\n",
    "        else:\n",
    "            conv_block = nn.Sequential(\n",
    "                nn.Conv2d(in_channels=channel[0], out_channels=channel[0], kernel_size=3, padding=1),\n",
    "                nn.Conv2d(in_channels=channel[0], out_channels=channel[1], kernel_size=1, padding=0),\n",
    "            )\n",
    "        return conv_block\n",
    "\n",
    "    def att_layer(self, channel):\n",
    "        att_block = nn.Sequential(\n",
    "            nn.Conv2d(in_channels=channel[0], out_channels=channel[1], kernel_size=1, padding=0),\n",
    "            nn.BatchNorm2d(channel[1]),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(in_channels=channel[1], out_channels=channel[2], kernel_size=1, padding=0),\n",
    "            nn.BatchNorm2d(channel[2]),\n",
    "            nn.Sigmoid(),\n",
    "        )\n",
    "        return att_block\n",
    "\n",
    "    def forward(self, x):\n",
    "        g_encoder, g_decoder, g_maxpool, g_upsampl, indices = ([0] * 5 for _ in range(5))\n",
    "        for i in range(5):\n",
    "            g_encoder[i], g_decoder[-i - 1] = ([0] * 2 for _ in range(2))\n",
    "\n",
    "        # define attention list for tasks\n",
    "        atten_encoder, atten_decoder = ([0] * 3 for _ in range(2))\n",
    "        for i in range(3):\n",
    "            atten_encoder[i], atten_decoder[i] = ([0] * 5 for _ in range(2))\n",
    "        for i in range(3):\n",
    "            for j in range(5):\n",
    "                atten_encoder[i][j], atten_decoder[i][j] = ([0] * 3 for _ in range(2))\n",
    "\n",
    "        # define global shared network\n",
    "        for i in range(5):\n",
    "            if i == 0:\n",
    "                g_encoder[i][0] = self.encoder_block[i](x)\n",
    "                g_encoder[i][1] = self.conv_block_enc[i](g_encoder[i][0])\n",
    "                g_maxpool[i], indices[i] = self.down_sampling(g_encoder[i][1])\n",
    "            else:\n",
    "                g_encoder[i][0] = self.encoder_block[i](g_maxpool[i - 1])\n",
    "                g_encoder[i][1] = self.conv_block_enc[i](g_encoder[i][0])\n",
    "                g_maxpool[i], indices[i] = self.down_sampling(g_encoder[i][1])\n",
    "\n",
    "        for i in range(5):\n",
    "            if i == 0:\n",
    "                g_upsampl[i] = self.up_sampling(g_maxpool[-1], indices[-i - 1])\n",
    "                g_decoder[i][0] = self.decoder_block[-i - 1](g_upsampl[i])\n",
    "                g_decoder[i][1] = self.conv_block_dec[-i - 1](g_decoder[i][0])\n",
    "            else:\n",
    "                g_upsampl[i] = self.up_sampling(g_decoder[i - 1][-1], indices[-i - 1])\n",
    "                g_decoder[i][0] = self.decoder_block[-i - 1](g_upsampl[i])\n",
    "                g_decoder[i][1] = self.conv_block_dec[-i - 1](g_decoder[i][0])\n",
    "\n",
    "        # define task dependent attention module\n",
    "        for i in range(3):\n",
    "            for j in range(5):\n",
    "                if j == 0:\n",
    "                    atten_encoder[i][j][0] = self.encoder_att[i][j](g_encoder[j][0])\n",
    "                    atten_encoder[i][j][1] = (atten_encoder[i][j][0]) * g_encoder[j][1]\n",
    "                    atten_encoder[i][j][2] = self.encoder_block_att[j](atten_encoder[i][j][1])\n",
    "                    atten_encoder[i][j][2] = F.max_pool2d(atten_encoder[i][j][2], kernel_size=2, stride=2)\n",
    "                else:\n",
    "                    atten_encoder[i][j][0] = self.encoder_att[i][j](torch.cat((g_encoder[j][0], atten_encoder[i][j - 1][2]), dim=1))\n",
    "                    atten_encoder[i][j][1] = (atten_encoder[i][j][0]) * g_encoder[j][1]\n",
    "                    atten_encoder[i][j][2] = self.encoder_block_att[j](atten_encoder[i][j][1])\n",
    "                    atten_encoder[i][j][2] = F.max_pool2d(atten_encoder[i][j][2], kernel_size=2, stride=2)\n",
    "        \n",
    "        return [att[-1][2] for att in atten_encoder] \n",
    "\n",
    "\n",
    "# MY EDIT\n",
    "#             for j in range(5):\n",
    "#                 if j == 0:\n",
    "#                     atten_decoder[i][j][0] = F.interpolate(atten_encoder[i][-1][-1], scale_factor=2, mode='bilinear', align_corners=True)\n",
    "#                     atten_decoder[i][j][0] = self.decoder_block_att[-j - 1](atten_decoder[i][j][0])\n",
    "#                     atten_decoder[i][j][1] = self.decoder_att[i][-j - 1](torch.cat((g_upsampl[j], atten_decoder[i][j][0]), dim=1))\n",
    "#                     atten_decoder[i][j][2] = (atten_decoder[i][j][1]) * g_decoder[j][-1]\n",
    "#                 else:\n",
    "#                     atten_decoder[i][j][0] = F.interpolate(atten_decoder[i][j - 1][2], scale_factor=2, mode='bilinear', align_corners=True)\n",
    "#                     atten_decoder[i][j][0] = self.decoder_block_att[-j - 1](atten_decoder[i][j][0])\n",
    "#                     atten_decoder[i][j][1] = self.decoder_att[i][-j - 1](torch.cat((g_upsampl[j], atten_decoder[i][j][0]), dim=1))\n",
    "#                     atten_decoder[i][j][2] = (atten_decoder[i][j][1]) * g_decoder[j][-1]\n",
    "\n",
    "#         # define task prediction layers\n",
    "#         t1_pred = F.log_softmax(self.pred_task1(atten_decoder[0][-1][-1]), dim=1)\n",
    "#         t2_pred = self.pred_task2(atten_decoder[1][-1][-1])\n",
    "#         t3_pred = self.pred_task3(atten_decoder[2][-1][-1])\n",
    "#         t3_pred = t3_pred / torch.norm(t3_pred, p=2, dim=1, keepdim=True)\n",
    "\n",
    "#         return [t1_pred, t2_pred, t3_pred], self.logsigma"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "loose-sector",
   "metadata": {},
   "source": [
    "### Test Wrapper\n",
    "- Wrapping my MTANVGG16 to output the encoder representations for each task, since we don't use the classifiers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "adult-portugal",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class TestVGG(MTANVGG16):\n",
    "    \"\"\"\n",
    "    Just early return representations before classifier\n",
    "    \"\"\"\n",
    "    def __init__(self):\n",
    "        super().__init__({str(i): i for i in range(3)}) # doesnt matter\n",
    "    \n",
    "    def forward(self, X):\n",
    "        sh_outs = []\n",
    "        for layer in self.shared_layers:\n",
    "            X = layer(X)\n",
    "            if isinstance(layer, nn.ReLU):\n",
    "                sh_outs.append(X)\n",
    "        \n",
    "        maxpool = lambda x: F.max_pool2d(x, kernel_size=2, stride=2)\n",
    "        sh_out_idx = [(0, 1), (2, 3), (4, 6), (7, 9), (10, 12)] # indices of the relevant outputs from the shared backbone that go into the attention module\n",
    "        \n",
    "        res = []\n",
    "        for task in [\"0\", \"1\", \"2\"]:\n",
    "            att_out = None\n",
    "            for block, (in1_idx, in2_idx) in enumerate(sh_out_idx):\n",
    "                att_out = maxpool(self.attentions[block](task, sh_outs[in1_idx], sh_outs[in2_idx], att_out))\n",
    "            res.append(att_out)\n",
    "        return res"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adult-station",
   "metadata": {},
   "source": [
    "### Initialisation\n",
    "- For the implementations to be able to output the same thing, they need to have the same network parameters\n",
    "- We only really care to copy the conv weights, everything else either doesnt have params or is initialised statically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "discrete-wrestling",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/mihai/work/libs/miniconda3/envs/dev/lib/python3.8/site-packages/torch/nn/init.py:405: UserWarning: Initializing zero-element tensors is a no-op\n",
      "  warnings.warn(\"Initializing zero-element tensors is a no-op\")\n"
     ]
    }
   ],
   "source": [
    "# init here to set the conv biases and BN layers\n",
    "my = init_model(TestVGG())\n",
    "ref = init_model(SegNet())\n",
    "\n",
    "# COPY conv weights\n",
    "with torch.no_grad():\n",
    "    # attentions\n",
    "    for block in range(5):\n",
    "        for task in range(3):\n",
    "            my.attentions[block].task_att[str(task)][0].weight.copy_(ref.encoder_att[task][block][0].weight)\n",
    "            my.attentions[block].task_att[str(task)][3].weight.copy_(ref.encoder_att[task][block][3].weight)\n",
    "        my.attentions[block].shared_att[0].weight.copy_(ref.encoder_block_att[block][0].weight)\n",
    "\n",
    "    # backbone\n",
    "    ref_backbone = nn.Sequential(*list(itertools.chain(*[list(itertools.chain(l1, l2)) for l1, l2 in zip(ref.encoder_block, ref.conv_block_enc)])))\n",
    "    ref_backbone_convs = [m for m in ref_backbone.modules() if isinstance(m, nn.Conv2d)]\n",
    "    my_backbone_convs = [m for m in my.shared_layers.modules() if isinstance(m, nn.Conv2d)]\n",
    "    for my_conv, ref_conv in zip(my_backbone_convs, ref_backbone_convs):\n",
    "        my_conv.weight.copy_(ref_conv.weight)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "integrated-opposition",
   "metadata": {},
   "source": [
    "### Comparing Outputs\n",
    "- Works :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "strategic-ukraine",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Outputs match for VGG16: True\n"
     ]
    }
   ],
   "source": [
    "X = torch.randn(1, 3, 256, 256)\n",
    "\n",
    "my_out = my(X)\n",
    "ref_out = ref(X)\n",
    "\n",
    "print(\"Outputs match for VGG16:\", all([(torch.equal(y1, y2)) for y1, y2 in zip(my_out, ref_out)]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "miniature-massachusetts",
   "metadata": {},
   "source": [
    "## ResNets\n",
    "- Not used for multi-domain classificaiton in paper\n",
    "- Usef as part of DeepLab architecture for dense predictions\n",
    "- Backbone can be changed to other resnets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "secret-moldova",
   "metadata": {},
   "source": [
    "### Original Implementation\n",
    "- We can comment out the Deeplab part and just keep the backbone\n",
    "- Orig implementation uses dilated resnets based on own implementation - would have had to copy paste those as well, so I just simplified and and using standard resnet as backbone\n",
    "- Adding maxpooling after stages to make up for not using dilation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "behind-basics",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MTANDeepLabv3(nn.Module):\n",
    "    \"\"\"\n",
    "    Orig Implementation from @https://github.com/lorenmt/mtan/blob/master/im2im_pred/model_resnet_mtan/resnet_mtan.py\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self):\n",
    "        super(MTANDeepLabv3, self).__init__()\n",
    "#         backbone = ResnetDilated(resnet.__dict__['resnet50'](pretrained=True)) # MY EDIT\n",
    "        backbone = resnet50() # Use default resnet instead of dilated for testing; just to minimise dependencies\n",
    "        ch = [256, 512, 1024, 2048]\n",
    "        \n",
    "        self.tasks = ['segmentation', 'depth', 'normal']\n",
    "        self.num_out_channels = {'segmentation': 13, 'depth': 1, 'normal': 3}\n",
    "        \n",
    "#         self.shared_conv = nn.Sequential(backbone.conv1, backbone.bn1, backbone.relu1, backbone.maxpool) # MY EDIT\n",
    "        self.shared_conv = nn.Sequential(backbone.conv1, backbone.bn1, backbone.relu, backbone.maxpool)\n",
    "\n",
    "        \n",
    "        \n",
    "        # We will apply the attention over the last bottleneck layer in the ResNet. \n",
    "        self.shared_layer1_b = backbone.layer1[:-1] \n",
    "        self.shared_layer1_t = backbone.layer1[-1]\n",
    "\n",
    "        self.shared_layer2_b = backbone.layer2[:-1]\n",
    "        self.shared_layer2_t = backbone.layer2[-1]\n",
    "\n",
    "        self.shared_layer3_b = backbone.layer3[:-1]\n",
    "        self.shared_layer3_t = backbone.layer3[-1]\n",
    "\n",
    "        self.shared_layer4_b = backbone.layer4[:-1]\n",
    "        self.shared_layer4_t = backbone.layer4[-1]\n",
    "\n",
    "        # Define task specific attention modules using a similar bottleneck design in residual block\n",
    "        # (to avoid large computations)\n",
    "        self.encoder_att_1 = nn.ModuleList([self.att_layer(ch[0], ch[0] // 4, ch[0]) for _ in self.tasks])\n",
    "        self.encoder_att_2 = nn.ModuleList([self.att_layer(2 * ch[1], ch[1] // 4, ch[1]) for _ in self.tasks])\n",
    "        self.encoder_att_3 = nn.ModuleList([self.att_layer(2 * ch[2], ch[2] // 4, ch[2]) for _ in self.tasks])\n",
    "        self.encoder_att_4 = nn.ModuleList([self.att_layer(2 * ch[3], ch[3] // 4, ch[3]) for _ in self.tasks])\n",
    "\n",
    "        # Define task shared attention encoders using residual bottleneck layers\n",
    "        # We do not apply shared attention encoders at the last layer,\n",
    "        # so the attended features will be directly fed into the task-specific decoders.\n",
    "        self.encoder_block_att_1 = self.conv_layer(ch[0], ch[1] // 4)\n",
    "        self.encoder_block_att_2 = self.conv_layer(ch[1], ch[2] // 4)\n",
    "        self.encoder_block_att_3 = self.conv_layer(ch[2], ch[3] // 4)\n",
    "        \n",
    "        self.down_sampling = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "\n",
    "        # Define task-specific decoders using ASPP modules\n",
    "#         self.decoders = nn.ModuleList([DeepLabHead(2048, self.num_out_channels[t]) for t in self.tasks])\n",
    "        \n",
    "    def forward(self, x, out_size=None):\n",
    "        # Shared convolution\n",
    "        x = self.shared_conv(x)\n",
    "        \n",
    "        # Shared ResNet block 1\n",
    "        u_1_b = self.shared_layer1_b(x)\n",
    "        u_1_t = self.shared_layer1_t(u_1_b)\n",
    "\n",
    "        # Shared ResNet block 2\n",
    "        u_2_b = self.shared_layer2_b(u_1_t)\n",
    "        u_2_t = self.shared_layer2_t(u_2_b)\n",
    "\n",
    "        # Shared ResNet block 3\n",
    "        u_3_b = self.shared_layer3_b(u_2_t)\n",
    "        u_3_t = self.shared_layer3_t(u_3_b)\n",
    "        \n",
    "        # Shared ResNet block 4\n",
    "        u_4_b = self.shared_layer4_b(u_3_t)\n",
    "        u_4_t = self.shared_layer4_t(u_4_b)\n",
    "\n",
    "        # Attention block 1 -> Apply attention over last residual block\n",
    "        a_1_mask = [att_i(u_1_b) for att_i in self.encoder_att_1]  # Generate task specific attention map\n",
    "        a_1 = [a_1_mask_i * u_1_t for a_1_mask_i in a_1_mask]  # Apply task specific attention map to shared features\n",
    "        a_1 = [self.down_sampling(self.encoder_block_att_1(a_1_i)) for a_1_i in a_1]\n",
    "    \n",
    "        # Attention block 2 -> Apply attention over last residual block\n",
    "        a_2_mask = [att_i(torch.cat((u_2_b, a_1_i), dim=1)) for a_1_i, att_i in zip(a_1, self.encoder_att_2)]\n",
    "        a_2 = [a_2_mask_i * u_2_t for a_2_mask_i in a_2_mask]\n",
    "#         a_2 = [self.encoder_block_att_2(a_2_i) for a_2_i in a_2] # MY EDIT\n",
    "        a_2 = [self.down_sampling(self.encoder_block_att_2(a_2_i)) for a_2_i in a_2] # NOTE: Add maxpooling (like in visual decathlon impl) becuse regular of resnet backbone\n",
    "        \n",
    "        # Attention block 3 -> Apply attention over last residual block\n",
    "        a_3_mask = [att_i(torch.cat((u_3_b, a_2_i), dim=1)) for a_2_i, att_i in zip(a_2, self.encoder_att_3)]\n",
    "        a_3 = [a_3_mask_i * u_3_t for a_3_mask_i in a_3_mask]\n",
    "#         a_3 = [self.encoder_block_att_3(a_3_i) for a_3_i in a_3] # MY EDIT\n",
    "        a_3 = [self.down_sampling(self.encoder_block_att_3(a_3_i)) for a_3_i in a_3] # NOTE: Add maxpooling (like in visual decathlon impl) becuse regular of resnet backbone\n",
    "        \n",
    "        # Attention block 4 -> Apply attention over last residual block (without final encoder)\n",
    "        a_4_mask = [att_i(torch.cat((u_4_b, a_3_i), dim=1)) for a_3_i, att_i in zip(a_3, self.encoder_att_4)]\n",
    "        a_4 = [a_4_mask_i * u_4_t for a_4_mask_i in a_4_mask]\n",
    "        \n",
    "        pred = [F.avg_pool2d(a_4_i, 8) for a_4_i in a_4]\n",
    "        return pred\n",
    "        \n",
    "#         # Task specific decoders\n",
    "#         out = [0 for _ in self.tasks]\n",
    "#         for i, t in enumerate(self.tasks):\n",
    "#             out[i] = F.interpolate(self.decoders[i](a_4[i]), size=out_size, mode='bilinear', align_corners=True)\n",
    "#             if t == 'segmentation':\n",
    "#                 out[i] = F.log_softmax(out[i], dim=1)\n",
    "#             if t == 'normal':\n",
    "#                 out[i] = out[i] / torch.norm(out[i], p=2, dim=1, keepdim=True)\n",
    "#         return out\n",
    "    \n",
    "    def att_layer(self, in_channel, intermediate_channel, out_channel):\n",
    "        return nn.Sequential(\n",
    "            nn.Conv2d(in_channels=in_channel, out_channels=intermediate_channel, kernel_size=1, padding=0),\n",
    "            nn.BatchNorm2d(intermediate_channel),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(in_channels=intermediate_channel, out_channels=out_channel, kernel_size=1, padding=0),\n",
    "            nn.BatchNorm2d(out_channel),\n",
    "            nn.Sigmoid())\n",
    "        \n",
    "    def conv_layer(self, in_channel, out_channel):\n",
    "        downsample = nn.Sequential(conv1x1(in_channel, 4 * out_channel, stride=1),\n",
    "                                   nn.BatchNorm2d(4 * out_channel))\n",
    "        return Bottleneck(in_channel, out_channel, downsample=downsample)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sporting-elite",
   "metadata": {},
   "source": [
    "### Test Wrapper\n",
    "- Wrapping my MTANResNet to output the encoder representations for each task, since we don't use the classifiers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "dynamic-requirement",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class TestResNet(MTANResNet):\n",
    "    def __init__(self):\n",
    "        super().__init__({str(i): i for i in range(3)}) # doesnt matter\n",
    "    \n",
    "    def forward(self, X):\n",
    "        out = self.shared_conv(X)\n",
    "        sh_outs = [out := l(out) for l in self.shared_layers]\n",
    "        \n",
    "        maxpool = lambda x: F.max_pool2d(x, kernel_size=2, stride=2)\n",
    "        \n",
    "        res = []\n",
    "        for task in [\"0\", \"1\", \"2\"]:\n",
    "            att_out = None\n",
    "            for stage in range(4):\n",
    "                att_out = self.attentions[stage](task, sh_outs[stage * 2], sh_outs[stage * 2 + 1], att_out)\n",
    "                if stage < 3:\n",
    "                    # NOTE: This can/should be changed if backbone is different (e.g uses dilation)\n",
    "                    att_out = maxpool(att_out)\n",
    "\n",
    "            out = F.avg_pool2d(att_out, 8)\n",
    "            res.append(out)\n",
    "        \n",
    "        return res"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "interior-bridge",
   "metadata": {},
   "source": [
    "### Initialisation\n",
    "- For the implementations to be able to output the same thing, they need to have the same network parameters\n",
    "- We only really care to copy the conv weights, everything else either doesnt have params or is initialised statically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "entertaining-programmer",
   "metadata": {},
   "outputs": [],
   "source": [
    "# init here to set the conv biases and BN layers\n",
    "my = init_model(TestResNet())\n",
    "ref = init_model(MTANDeepLabv3())\n",
    "\n",
    "# COPY conv weights\n",
    "with torch.no_grad():\n",
    "    # attentions\n",
    "    for stage in range(4):\n",
    "        my_att = my.attentions[stage]\n",
    "        \n",
    "        # task specific\n",
    "        for task in range(3):\n",
    "            my_att.task_att[str(task)][0].weight.copy_(getattr(ref, f\"encoder_att_{stage+1}\")[task][0].weight)\n",
    "            my_att.task_att[str(task)][3].weight.copy_(getattr(ref, f\"encoder_att_{stage+1}\")[task][3].weight)\n",
    "        \n",
    "        # shared\n",
    "        if len(list(my_att.shared_att.modules())) > 1:\n",
    "            ref_convs = [m for m in getattr(ref, f\"encoder_block_att_{stage+1}\").modules() if isinstance(m, nn.Conv2d)]\n",
    "            my_convs = [m for m in my_att.shared_att.modules() if isinstance(m, nn.Conv2d)]\n",
    "            for my_conv, ref_conv in zip(my_convs, ref_convs):\n",
    "                my_conv.weight.copy_(ref_conv.weight)\n",
    "\n",
    "    # backbone\n",
    "    my.shared_conv[0].weight.copy_(ref.shared_conv[0].weight)\n",
    "    ref_backbone = [f\"shared_layer{stage}_{split}\" for stage in range(1, 5) for split in [\"b\", \"t\"]]\n",
    "    ref_backbone = nn.Sequential(*[getattr(ref, layer) for layer in ref_backbone])\n",
    "    ref_backbone_convs = [m for m in ref_backbone.modules() if isinstance(m, nn.Conv2d)]\n",
    "    my_backbone_convs = [m for m in my.shared_layers.modules() if isinstance(m, nn.Conv2d)]\n",
    "    for my_conv, ref_conv in zip(my_backbone_convs, ref_backbone_convs):\n",
    "        my_conv.weight.copy_(ref_conv.weight)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "framed-bahamas",
   "metadata": {},
   "source": [
    "### Comparing Outputs\n",
    "- Works :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "japanese-little",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Outputs match for ResNet50: True\n"
     ]
    }
   ],
   "source": [
    "X = torch.randn(1, 3, 256, 256)\n",
    "\n",
    "my_out = my(X)\n",
    "ref_out = ref(X)\n",
    "\n",
    "print(\"Outputs match for ResNet50:\", all([torch.equal(y1, y2) for y1, y2 in zip(my_out, ref_out)]))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:dev]",
   "language": "python",
   "name": "conda-env-dev-py"
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}