raw_yolact.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/broadinstitute/superurop-log/blob/main/raw_yolact.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h2CKA97pOL8t"
      },
      "outputs": [],
      "source": [
        "import torchvision\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nO1g_QEjO0Uf",
        "outputId": "5f72a633-4e68-4c3d-b98c-943cc7475c27"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Downloading: \"https://download.pytorch.org/models/resnet101-cd907fc2.pth\" to /root/.cache/torch/hub/checkpoints/resnet101-cd907fc2.pth\n",
            "100%|██████████| 171M/171M [00:02<00:00, 77.7MB/s]\n"
          ]
        }
      ],
      "source": [
        "res_weights = torchvision.models.ResNet101_Weights.DEFAULT\n",
        "resnet101 = torchvision.models.resnet101(weights=res_weights, progress=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dCtDEFyNTlvK",
        "outputId": "bfd13671-e0e3-41e9-c16b-01029bad25aa"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "torch.Size([1, 2048, 18, 18])"
            ]
          },
          "execution_count": 4,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "layers = list(resnet101.children())[:-2]\n",
        "net = torch.nn.Sequential(*layers)\n",
        "net.eval()\n",
        "net(torch.randn(1, 3, 550, 550)).shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "33AosadrwybG"
      },
      "outputs": [],
      "source": [
        "from collections import OrderedDict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qd1Vn-y1uSfO"
      },
      "outputs": [],
      "source": [
        "class Bottleneck(nn.Module):\n",
        "    \"\"\" Adapted from torchvision.models.resnet \"\"\"\n",
        "    expansion = 4\n",
        "\n",
        "    def __init__(self, inplanes, planes, stride=1, downsample=None, norm_layer=nn.BatchNorm2d, dilation=1):\n",
        "        super(Bottleneck, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False, dilation=dilation)\n",
        "        self.bn1 = norm_layer(planes)\n",
        "        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,\n",
        "                                padding=dilation, bias=False, dilation=dilation)\n",
        "        self.bn2 = norm_layer(planes)\n",
        "        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False, dilation=dilation)\n",
        "        self.bn3 = norm_layer(planes * 4)\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "        self.downsample = downsample\n",
        "        self.stride = stride\n",
        "\n",
        "    def forward(self, x):\n",
        "        residual = x\n",
        "\n",
        "        out = self.conv1(x)\n",
        "        out = self.bn1(out)\n",
        "        out = self.relu(out)\n",
        "\n",
        "        out = self.conv2(out)\n",
        "        out = self.bn2(out)\n",
        "        out = self.relu(out)\n",
        "\n",
        "        out = self.conv3(out)\n",
        "        out = self.bn3(out)\n",
        "\n",
        "        if self.downsample is not None:\n",
        "            residual = self.downsample(x)\n",
        "\n",
        "        out += residual\n",
        "        out = self.relu(out)\n",
        "\n",
        "        return out"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-QOLPHp2ufSS",
        "outputId": "7394c9c5-3e62-4ce7-f482-fbd3b1deb11e"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Bottleneck(\n",
              "  (conv1): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "  (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "  (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "  (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "  (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "  (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "  (relu): ReLU(inplace=True)\n",
              ")"
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "bn = Bottleneck(512, 512)\n",
        "\n",
        "bn"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Rm7r3kokvLqe"
      },
      "outputs": [],
      "source": [
        "class ResNetBackbone(nn.Module):\n",
        "    \"\"\" Adapted from torchvision.models.resnet \"\"\"\n",
        "\n",
        "    def __init__(self, layers, atrous_layers=[], block=Bottleneck, norm_layer=nn.BatchNorm2d):\n",
        "        super().__init__()\n",
        "\n",
        "        # These will be populated by _make_layer\n",
        "        self.num_base_layers = len(layers)\n",
        "        self.layers = nn.ModuleList()\n",
        "        self.channels = []\n",
        "        self.norm_layer = norm_layer\n",
        "        self.dilation = 1\n",
        "        self.atrous_layers = atrous_layers\n",
        "\n",
        "        # From torchvision.models.resnet.Resnet\n",
        "        self.inplanes = 64\n",
        "\n",
        "        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)\n",
        "        self.bn1 = norm_layer(64)\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
        "\n",
        "        self._make_layer(block, 64, layers[0])\n",
        "        self._make_layer(block, 128, layers[1], stride=2)\n",
        "        self._make_layer(block, 256, layers[2], stride=2)\n",
        "        self._make_layer(block, 512, layers[3], stride=2)\n",
        "\n",
        "        # This contains every module that should be initialized by loading in pretrained weights.\n",
        "        # Any extra layers added onto this that won't be initialized by init_backbone will not be\n",
        "        # in this list. That way, Yolact::init_weights knows which backbone weights to initialize\n",
        "        # with xavier, and which ones to leave alone.\n",
        "        self.backbone_modules = [m for m in self.modules() if isinstance(m, nn.Conv2d)]\n",
        "\n",
        "\n",
        "    def _make_layer(self, block, planes, blocks, stride=1):\n",
        "        \"\"\" Here one layer means a string of n Bottleneck blocks. \"\"\"\n",
        "        downsample = None\n",
        "\n",
        "        if stride != 1 or self.inplanes != planes * block.expansion:\n",
        "            if len(self.layers) in self.atrous_layers:\n",
        "                self.dilation += 1\n",
        "                stride = 1\n",
        "\n",
        "            downsample = nn.Sequential(\n",
        "                nn.Conv2d(self.inplanes, planes * block.expansion,\n",
        "                          kernel_size=1, stride=stride, bias=False,\n",
        "                          dilation=self.dilation),\n",
        "                self.norm_layer(planes * block.expansion),\n",
        "            )\n",
        "\n",
        "        layers = []\n",
        "        layers.append(block(self.inplanes, planes, stride, downsample, self.norm_layer, self.dilation))\n",
        "        self.inplanes = planes * block.expansion\n",
        "        for i in range(1, blocks):\n",
        "            layers.append(block(self.inplanes, planes, norm_layer=self.norm_layer))\n",
        "        layer = nn.Sequential(*layers)\n",
        "\n",
        "        self.channels.append(planes * block.expansion)\n",
        "        self.layers.append(layer)\n",
        "\n",
        "        return layer\n",
        "\n",
        "    def init_weights(self):  # TODO: match key names\n",
        "        # resnet101 pretrained weights\n",
        "\n",
        "        # state_dict = torch.hub.load_state_dict_from_url(\n",
        "        #     'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',\n",
        "        #     progress=True)\n",
        "        # self.load_state_dict(state_dict)\n",
        "\n",
        "        model = torchvision.models.resnet101(pretrained=True)\n",
        "        state_dict = {k:v for k, v in model.state_dict().items() if k in self.state_dict()}\n",
        "        self.state_dict = state_dict\n",
        "        # self.load_state_dict(state_dict)\n",
        "\n",
        "    def forward(self, x):\n",
        "        \"\"\" Returns a list of convouts for each layer. \"\"\"\n",
        "\n",
        "        x = self.conv1(x)\n",
        "        x = self.bn1(x)\n",
        "        x = self.relu(x)\n",
        "        x = self.maxpool(x)\n",
        "\n",
        "        outs = []\n",
        "        for layer in self.layers:\n",
        "            x = layer(x)\n",
        "            outs.append(x)\n",
        "\n",
        "        return tuple(outs)\n",
        "\n",
        "    def add_layer(self, conv_channels=1024, downsample=2, depth=1, block=Bottleneck):\n",
        "        \"\"\" Add a downsample layer to the backbone as per what SSD does. \"\"\"\n",
        "        self._make_layer(block, conv_channels // block.expansion, blocks=depth, stride=downsample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xkyb84crxw_f",
        "outputId": "1fb29056-82c8-4a3d-c54a-03ae54afed82"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
            "  warnings.warn(\n",
            "/usr/local/lib/python3.10/dist-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=ResNet101_Weights.IMAGENET1K_V1`. You can also use `weights=ResNet101_Weights.DEFAULT` to get the most up-to-date weights.\n",
            "  warnings.warn(msg)\n",
            "Downloading: \"https://download.pytorch.org/models/resnet101-63fe2227.pth\" to /root/.cache/torch/hub/checkpoints/resnet101-63fe2227.pth\n",
            "100%|██████████| 171M/171M [00:01<00:00, 124MB/s]\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "ResNetBackbone(\n",
              "  (layers): ModuleList(\n",
              "    (0): Sequential(\n",
              "      (0): Bottleneck(\n",
              "        (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "        (downsample): Sequential(\n",
              "          (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        )\n",
              "      )\n",
              "      (1): Bottleneck(\n",
              "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (2): Bottleneck(\n",
              "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "    )\n",
              "    (1): Sequential(\n",
              "      (0): Bottleneck(\n",
              "        (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "        (downsample): Sequential(\n",
              "          (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
              "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        )\n",
              "      )\n",
              "      (1): Bottleneck(\n",
              "        (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (2): Bottleneck(\n",
              "        (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (3): Bottleneck(\n",
              "        (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "    )\n",
              "    (2): Sequential(\n",
              "      (0): Bottleneck(\n",
              "        (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "        (downsample): Sequential(\n",
              "          (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
              "          (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        )\n",
              "      )\n",
              "      (1): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (2): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (3): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (4): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (5): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (6): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (7): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (8): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (9): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (10): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (11): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (12): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (13): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (14): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (15): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (16): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (17): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (18): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (19): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (20): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (21): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (22): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "    )\n",
              "    (3): Sequential(\n",
              "      (0): Bottleneck(\n",
              "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "        (downsample): Sequential(\n",
              "          (0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
              "          (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        )\n",
              "      )\n",
              "      (1): Bottleneck(\n",
              "        (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "      (2): Bottleneck(\n",
              "        (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
              "        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "      )\n",
              "    )\n",
              "  )\n",
              "  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
              "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "  (relu): ReLU(inplace=True)\n",
              "  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
              ")"
            ]
          },
          "execution_count": 17,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "bn = ResNetBackbone([3, 4, 23, 3])\n",
        "bn.init_weights()\n",
        "bn"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Vy_f5SRb5I2V",
        "outputId": "b8f96be5-0c17-4f03-d623-92f230066e5f"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'conv1.weight': tensor([[[[ 2.0222e-02, -4.3927e-03, -1.8274e-02,  ..., -1.5180e-02,\n",
              "            -1.5794e-03,  9.3115e-03],\n",
              "           [-4.0752e-03,  3.2116e-03, -1.5956e-02,  ..., -8.4465e-02,\n",
              "            -7.4997e-02, -4.0676e-02],\n",
              "           [ 3.5039e-03,  2.6746e-02,  5.0813e-02,  ...,  3.3407e-02,\n",
              "             1.3659e-02,  2.7821e-02],\n",
              "           ...,\n",
              "           [-3.6174e-02, -1.2986e-01, -3.0369e-01,  ..., -3.7412e-01,\n",
              "            -1.3025e-01,  4.2633e-02],\n",
              "           [ 1.5479e-02,  2.3444e-02,  6.5222e-03,  ..., -1.6439e-01,\n",
              "            -1.8245e-01, -9.7434e-02],\n",
              "           [-3.0444e-02, -1.1357e-02,  4.9984e-02,  ...,  1.6412e-01,\n",
              "             1.0419e-01, -1.2681e-02]],\n",
              " \n",
              "          [[ 8.7115e-03, -5.8911e-03, -1.2204e-02,  ..., -1.3515e-02,\n",
              "             1.5212e-02,  1.9115e-02],\n",
              "           [-6.8970e-03,  1.0470e-02, -7.7561e-03,  ..., -7.9215e-02,\n",
              "            -5.9150e-02, -2.1380e-02],\n",
              "           [-2.4955e-03,  3.2179e-02,  7.6542e-02,  ...,  8.9056e-02,\n",
              "             5.1445e-02,  3.4868e-02],\n",
              "           ...,\n",
              "           [-6.3695e-02, -1.9335e-01, -4.2540e-01,  ..., -5.4060e-01,\n",
              "            -2.1634e-01,  8.1145e-03],\n",
              "           [-9.9660e-03,  9.2712e-04, -4.3920e-02,  ..., -2.9635e-01,\n",
              "            -2.8675e-01, -1.7962e-01],\n",
              "           [ 2.5959e-02,  7.8301e-02,  1.5091e-01,  ...,  2.7260e-01,\n",
              "             2.0840e-01,  4.2282e-02]],\n",
              " \n",
              "          [[ 4.0026e-03, -9.1936e-03, -2.0569e-02,  ...,  1.0229e-03,\n",
              "             1.5326e-02,  9.7999e-04],\n",
              "           [ 1.7419e-02,  3.0881e-02,  6.3634e-03,  ..., -1.2466e-02,\n",
              "             7.4009e-03,  2.0044e-02],\n",
              "           [-1.3397e-02, -5.1027e-03, -3.2198e-02,  ..., -5.9151e-02,\n",
              "            -3.1463e-02, -9.3953e-03],\n",
              "           ...,\n",
              "           [-2.4370e-02, -8.1988e-02, -1.8747e-01,  ..., -2.2849e-01,\n",
              "            -4.5930e-02,  5.7061e-02],\n",
              "           [-4.6034e-03,  1.7281e-02,  2.0939e-02,  ..., -7.0717e-02,\n",
              "            -1.0450e-01, -8.8036e-02],\n",
              "           [-6.7201e-04,  1.1773e-02,  3.1304e-02,  ...,  1.1040e-01,\n",
              "             9.6085e-02, -5.0325e-03]]],\n",
              " \n",
              " \n",
              "         [[[-1.3283e-02, -6.3476e-03,  1.7371e-02,  ..., -2.6374e-02,\n",
              "            -3.9297e-02,  3.6769e-02],\n",
              "           [ 1.6402e-02, -9.2402e-03,  6.0237e-03,  ...,  4.5398e-02,\n",
              "             6.1903e-02, -7.1906e-02],\n",
              "           [ 1.0526e-02,  1.2171e-02,  9.3159e-02,  ..., -3.4628e-01,\n",
              "             6.6404e-02,  7.0177e-02],\n",
              "           ...,\n",
              "           [-3.6881e-02,  7.5082e-02, -1.3295e-01,  ...,  5.8651e-01,\n",
              "             1.0566e-01, -1.2248e-01],\n",
              "           [ 2.1673e-02, -5.9233e-02,  2.3098e-01,  ..., -2.1465e-01,\n",
              "             2.1108e-01, -1.8771e-02],\n",
              "           [ 6.6413e-03, -5.1895e-02, -1.2541e-02,  ..., -8.8446e-02,\n",
              "            -1.8269e-02,  3.2724e-02]],\n",
              " \n",
              "          [[-5.7876e-03, -6.5713e-03,  9.7605e-03,  ..., -2.4045e-02,\n",
              "            -4.0857e-02,  3.7269e-02],\n",
              "           [ 1.7980e-02, -2.9323e-02,  9.4012e-03,  ...,  3.5879e-02,\n",
              "             7.4789e-02, -3.7304e-02],\n",
              "           [ 1.8119e-02,  2.5782e-03,  9.7156e-02,  ..., -4.4105e-01,\n",
              "             1.4873e-02,  1.0465e-01],\n",
              "           ...,\n",
              "           [-4.4777e-02,  8.7919e-02, -2.3434e-01,  ...,  7.4114e-01,\n",
              "             1.5556e-01, -1.4555e-01],\n",
              "           [ 3.1682e-02, -1.3436e-02,  2.4753e-01,  ..., -2.5792e-01,\n",
              "             2.6018e-01, -1.2641e-02],\n",
              "           [ 2.1700e-02, -8.5579e-04,  2.5962e-02,  ..., -1.0650e-01,\n",
              "            -2.9765e-02,  1.8971e-03]],\n",
              " \n",
              "          [[-5.0954e-03,  1.0142e-02,  2.8154e-03,  ..., -2.2294e-03,\n",
              "            -3.5401e-02,  1.8325e-02],\n",
              "           [ 6.3929e-03, -1.6070e-02, -1.0550e-02,  ...,  9.2290e-02,\n",
              "             6.2968e-02, -7.0256e-02],\n",
              "           [-3.6175e-03, -6.6921e-03,  5.0055e-02,  ..., -2.8415e-01,\n",
              "             1.3017e-01,  7.0130e-02],\n",
              "           ...,\n",
              "           [-3.1363e-02,  7.8047e-02, -7.0873e-02,  ...,  4.2964e-01,\n",
              "             3.5380e-02, -8.7354e-02],\n",
              "           [ 3.0397e-02, -5.8957e-02,  2.3035e-01,  ..., -2.0295e-01,\n",
              "             1.3268e-01, -1.0068e-02],\n",
              "           [ 8.1602e-03, -2.7920e-02, -3.8539e-02,  ..., -2.7914e-02,\n",
              "            -2.1438e-02,  8.2988e-03]]],\n",
              " \n",
              " \n",
              "         [[[ 5.9179e-03,  1.1450e-02, -6.1776e-02,  ..., -4.6299e-02,\n",
              "             6.7527e-02,  9.8275e-03],\n",
              "           [ 8.8669e-03,  7.3341e-02, -3.2901e-02,  ..., -1.0453e-01,\n",
              "             1.7328e-01,  4.7307e-02],\n",
              "           [-1.1470e-02,  1.2149e-01,  2.7013e-03,  ..., -1.5456e-01,\n",
              "             2.4886e-01,  3.5192e-02],\n",
              "           ...,\n",
              "           [ 1.2814e-02,  1.1176e-01, -1.3131e-02,  ..., -8.0056e-02,\n",
              "             2.6456e-01,  4.8966e-02],\n",
              "           [ 9.1849e-03,  8.3514e-02, -3.2001e-02,  ..., -4.2758e-02,\n",
              "             1.8086e-01,  2.9671e-02],\n",
              "           [-3.6703e-03,  8.5900e-03, -7.5705e-02,  ..., -3.2137e-02,\n",
              "             1.1875e-01,  1.3688e-02]],\n",
              " \n",
              "          [[-3.1539e-03,  3.1529e-02, -2.7158e-02,  ..., -7.4867e-02,\n",
              "             4.2428e-02,  1.5711e-02],\n",
              "           [ 2.8099e-02,  1.2275e-01, -5.9815e-03,  ..., -1.7016e-01,\n",
              "             1.5623e-01,  7.8585e-02],\n",
              "           [ 2.5077e-02,  2.0189e-01,  3.0908e-02,  ..., -2.2413e-01,\n",
              "             2.7603e-01,  9.4537e-02],\n",
              "           ...,\n",
              "           [ 3.5136e-02,  1.7748e-01, -1.4339e-02,  ..., -1.4229e-01,\n",
              "             2.9657e-01,  1.0562e-01],\n",
              "           [ 1.0093e-02,  9.9101e-02, -4.6115e-02,  ..., -8.1636e-02,\n",
              "             1.8805e-01,  6.2061e-02],\n",
              "           [ 8.4642e-03,  3.8329e-02, -5.6497e-02,  ..., -5.7700e-02,\n",
              "             1.1272e-01,  4.3174e-02]],\n",
              " \n",
              "          [[-1.4282e-04, -2.3389e-03, -1.9654e-02,  ..., -2.2358e-02,\n",
              "             1.9954e-02, -3.2502e-02],\n",
              "           [-2.6001e-03,  5.5482e-02,  1.5250e-02,  ..., -8.4497e-02,\n",
              "             6.9498e-02, -1.2412e-02],\n",
              "           [-2.2391e-02,  9.0480e-02,  6.0876e-02,  ..., -1.0753e-01,\n",
              "             1.3897e-01, -4.1540e-02],\n",
              "           ...,\n",
              "           [-1.2348e-02,  8.0637e-02,  5.6340e-02,  ..., -4.5426e-02,\n",
              "             1.3194e-01, -4.2526e-02],\n",
              "           [-1.7389e-02,  4.7423e-02,  1.9719e-02,  ..., -2.9090e-02,\n",
              "             8.3996e-02, -2.8975e-02],\n",
              "           [-8.7350e-03,  2.0953e-02, -4.2252e-03,  ..., -2.5390e-02,\n",
              "             4.4542e-02, -3.0590e-02]]],\n",
              " \n",
              " \n",
              "         ...,\n",
              " \n",
              " \n",
              "         [[[ 2.1164e-02,  1.6922e-02, -2.4637e-02,  ..., -3.5926e-03,\n",
              "            -4.5143e-03, -4.7802e-03],\n",
              "           [ 1.5557e-02,  4.1134e-02,  9.2594e-03,  ...,  6.8536e-02,\n",
              "            -3.3796e-02, -1.2293e-01],\n",
              "           [-2.9705e-02,  1.1398e-02, -1.8864e-02,  ...,  1.9850e-01,\n",
              "             1.9440e-01, -1.2550e-01],\n",
              "           ...,\n",
              "           [ 1.4001e-02, -4.0922e-02,  2.5370e-01,  ..., -4.6508e-01,\n",
              "            -1.2693e-01,  1.3241e-01],\n",
              "           [-5.6044e-03, -1.0473e-01, -1.6365e-02,  ...,  4.1196e-02,\n",
              "            -1.8936e-01,  2.9299e-03],\n",
              "           [ 1.6253e-02, -5.6961e-02, -1.7359e-01,  ...,  1.4994e-01,\n",
              "            -4.3681e-03,  3.1854e-02]],\n",
              " \n",
              "          [[-9.8924e-03,  5.4293e-03, -2.6138e-02,  ...,  2.9590e-02,\n",
              "             7.4812e-03, -2.1530e-02],\n",
              "           [-1.5667e-02,  1.3143e-02, -1.7931e-02,  ...,  1.3352e-01,\n",
              "             4.8710e-02, -5.9767e-02],\n",
              "           [-4.3947e-02, -1.8217e-02, -9.7404e-02,  ...,  2.1595e-01,\n",
              "             3.1335e-01, -1.8402e-03],\n",
              "           ...,\n",
              "           [ 6.3135e-02,  2.6133e-02,  3.3600e-01,  ..., -7.0122e-01,\n",
              "            -2.6651e-01,  1.6588e-01],\n",
              "           [ 7.3457e-03, -5.4677e-02,  8.6192e-02,  ..., -1.7732e-02,\n",
              "            -3.2929e-01, -5.9486e-02],\n",
              "           [ 3.2153e-03, -3.1568e-02, -9.5821e-02,  ...,  1.8495e-01,\n",
              "            -7.3682e-02, -3.6894e-02]],\n",
              " \n",
              "          [[-5.5842e-03, -2.7447e-03, -2.5824e-02,  ...,  9.0967e-04,\n",
              "             1.8520e-02, -2.8618e-03],\n",
              "           [-1.3200e-02,  5.1415e-03, -1.6833e-02,  ...,  3.9462e-02,\n",
              "             1.5885e-03, -5.6381e-02],\n",
              "           [-2.9015e-02, -5.7217e-03, -3.4202e-02,  ...,  1.6668e-01,\n",
              "             1.6170e-01, -6.7241e-02],\n",
              "           ...,\n",
              "           [ 3.5389e-02, -2.1586e-02,  1.8604e-01,  ..., -4.1291e-01,\n",
              "            -1.1603e-01,  1.0927e-01],\n",
              "           [ 2.3013e-03, -4.1210e-02, -1.8644e-02,  ...,  2.2626e-02,\n",
              "            -2.0080e-01, -3.4115e-02],\n",
              "           [ 1.9892e-02, -1.5134e-03, -1.0327e-01,  ...,  1.1701e-01,\n",
              "            -3.9063e-02, -1.6900e-02]]],\n",
              " \n",
              " \n",
              "         [[[-5.5538e-08, -2.6325e-08, -6.9138e-10,  ...,  9.7495e-09,\n",
              "             1.3223e-08,  1.8906e-08],\n",
              "           [-2.6514e-08, -1.5337e-08, -1.8666e-08,  ..., -2.9065e-08,\n",
              "             5.2469e-09,  1.9422e-08],\n",
              "           [-5.8648e-09, -1.9407e-08, -1.7830e-08,  ..., -4.9517e-08,\n",
              "            -4.3926e-08,  5.7348e-09],\n",
              "           ...,\n",
              "           [-7.6421e-08, -8.3873e-08, -6.7820e-08,  ..., -3.4331e-08,\n",
              "            -2.4019e-08,  1.3504e-09],\n",
              "           [-5.3446e-08, -6.0856e-08, -5.9789e-08,  ...,  7.6961e-09,\n",
              "            -7.3536e-09,  9.9825e-09],\n",
              "           [-7.5073e-08, -5.5320e-08, -4.5796e-08,  ...,  9.2459e-09,\n",
              "            -3.1809e-09,  8.4687e-09]],\n",
              " \n",
              "          [[ 1.1481e-08,  4.6150e-08,  7.6648e-08,  ...,  8.4643e-08,\n",
              "             8.0418e-08,  7.5105e-08],\n",
              "           [ 3.4774e-08,  5.7570e-08,  6.4259e-08,  ...,  5.5634e-08,\n",
              "             8.3158e-08,  8.7862e-08],\n",
              "           [ 5.0421e-08,  3.9131e-08,  4.5809e-08,  ...,  3.1533e-08,\n",
              "             4.0924e-08,  7.2495e-08],\n",
              "           ...,\n",
              "           [-2.5365e-08, -2.9820e-08, -1.2614e-08,  ...,  4.6912e-08,\n",
              "             5.6400e-08,  7.3814e-08],\n",
              "           [-2.6964e-09, -3.6669e-10,  6.5934e-10,  ...,  8.3020e-08,\n",
              "             6.4033e-08,  7.0015e-08],\n",
              "           [-1.9816e-08,  4.0452e-09,  9.7432e-09,  ...,  7.3291e-08,\n",
              "             4.6672e-08,  6.1159e-08]],\n",
              " \n",
              "          [[-2.2916e-08,  9.8197e-09,  4.0514e-08,  ...,  4.4516e-08,\n",
              "             3.9299e-08,  3.6953e-08],\n",
              "           [ 1.9321e-08,  2.4614e-08,  2.8277e-08,  ...,  2.2826e-08,\n",
              "             4.8924e-08,  4.8514e-08],\n",
              "           [ 6.3043e-08,  3.3363e-08,  3.0949e-08,  ...,  9.2149e-09,\n",
              "             1.1062e-08,  3.2589e-08],\n",
              "           ...,\n",
              "           [-1.7967e-08, -2.9080e-08, -7.1967e-10,  ...,  3.7657e-08,\n",
              "             3.3544e-08,  2.4920e-08],\n",
              "           [ 5.6783e-09,  3.1829e-09,  5.8844e-09,  ...,  6.8484e-08,\n",
              "             2.9343e-08,  1.6371e-08],\n",
              "           [-2.8144e-09,  1.3160e-08,  1.4706e-08,  ...,  4.7713e-08,\n",
              "             1.7250e-08,  2.1211e-08]]],\n",
              " \n",
              " \n",
              "         [[[-6.7577e-03, -2.2072e-02, -1.6120e-02,  ...,  1.9965e-02,\n",
              "            -2.3618e-02, -4.0877e-02],\n",
              "           [-1.0744e-02,  1.0183e-04,  4.3173e-03,  ...,  3.2774e-02,\n",
              "             8.8146e-03, -1.4997e-02],\n",
              "           [-1.8087e-02, -5.0862e-03,  1.9248e-02,  ...,  7.9772e-02,\n",
              "             5.7868e-02,  4.2854e-02],\n",
              "           ...,\n",
              "           [ 1.0578e-02,  4.2848e-02,  7.4620e-02,  ...,  1.1236e-01,\n",
              "             1.0815e-01,  1.0924e-01],\n",
              "           [-3.2028e-02,  3.7799e-03,  5.1643e-02,  ...,  9.1239e-02,\n",
              "             6.8991e-02,  4.9936e-02],\n",
              "           [-5.3535e-02, -1.3198e-02,  1.6904e-02,  ...,  8.6864e-02,\n",
              "             4.1772e-02,  3.9772e-02]],\n",
              " \n",
              "          [[-2.0383e-02, -3.3167e-02, -2.6665e-02,  ...,  3.5744e-04,\n",
              "            -2.6331e-02, -4.0418e-02],\n",
              "           [-1.3169e-02,  1.9657e-03,  1.1509e-02,  ...,  3.5650e-02,\n",
              "             1.7746e-02, -9.1710e-03],\n",
              "           [-2.7688e-02,  5.5393e-03,  3.6324e-02,  ...,  7.0741e-02,\n",
              "             5.0612e-02,  2.8343e-02],\n",
              "           ...,\n",
              "           [ 1.1707e-02,  4.1271e-02,  7.3833e-02,  ...,  9.3581e-02,\n",
              "             1.0225e-01,  9.0626e-02],\n",
              "           [-1.0045e-02,  1.7943e-02,  5.3709e-02,  ...,  8.3323e-02,\n",
              "             7.8453e-02,  4.8212e-02],\n",
              "           [-2.6048e-02, -1.3092e-04,  1.8727e-02,  ...,  6.5370e-02,\n",
              "             2.8696e-02,  2.4652e-02]],\n",
              " \n",
              "          [[ 3.3157e-02,  1.2283e-03,  1.3401e-03,  ...,  2.2791e-02,\n",
              "            -2.8388e-02, -5.0342e-02],\n",
              "           [ 2.1755e-02,  8.9925e-03,  6.1698e-03,  ...,  1.7188e-02,\n",
              "            -1.5088e-02, -4.6674e-02],\n",
              "           [-4.7259e-03, -3.3812e-03,  5.7801e-03,  ...,  2.7695e-02,\n",
              "             1.3164e-02, -1.7174e-02],\n",
              "           ...,\n",
              "           [ 2.3359e-03, -1.1652e-02, -4.5875e-03,  ..., -8.4557e-03,\n",
              "             1.5695e-02, -2.3720e-04],\n",
              "           [-6.0332e-03, -2.9486e-02, -9.2786e-03,  ...,  1.7190e-02,\n",
              "             7.1863e-03, -2.3858e-02],\n",
              "           [-6.8612e-03, -3.4898e-02, -3.2312e-02,  ..., -1.0367e-02,\n",
              "            -3.9500e-02, -3.4689e-02]]]]),\n",
              " 'bn1.weight': tensor([ 2.6052e-01,  1.9547e-01,  2.7345e-01,  4.2496e-01,  2.8262e-01,\n",
              "          1.5043e-08,  2.1773e-01,  4.8107e-09,  2.3715e-08,  1.7285e-01,\n",
              "          2.5586e-01,  3.1999e-01,  2.1336e-01,  3.9580e-01,  2.4105e-01,\n",
              "          3.1630e-01,  3.2008e-01,  1.2258e-08,  2.5178e-01,  4.1141e-01,\n",
              "          3.4347e-01,  2.5406e-01,  3.6061e-01,  2.2518e-01,  2.2779e-01,\n",
              "          3.9869e-01,  2.3415e-01,  1.8959e-08,  2.8683e-01,  2.4818e-01,\n",
              "          2.8782e-01,  2.7092e-01,  3.5531e-01,  2.5614e-01,  2.2668e-01,\n",
              "          2.0158e-01,  3.2884e-01,  1.1636e-08,  2.0610e-01,  2.4417e-01,\n",
              "          2.2894e-01,  2.4786e-08,  6.8058e-09,  9.0608e-09,  2.3953e-01,\n",
              "          2.5992e-01,  1.6803e-03,  2.3422e-01, -7.3777e-09,  3.5465e-01,\n",
              "          2.5320e-08,  3.8923e-01,  2.7656e-01,  1.4901e-08,  2.3883e-01,\n",
              "          1.5257e-08,  3.0032e-01,  2.0780e-01,  2.9937e-01,  2.5019e-01,\n",
              "          3.6788e-01,  2.2675e-01, -5.0446e-09,  4.2316e-01]),\n",
              " 'bn1.bias': tensor([ 1.9693e-01,  2.2315e-01,  1.8429e-01,  1.0665e+00,  6.8113e-01,\n",
              "         -5.1364e-08,  2.1928e-01, -2.8729e-08, -7.2996e-08, -1.6934e-02,\n",
              "          1.2930e-01,  6.7223e-01,  2.1728e-01, -3.7019e-01, -1.0360e-01,\n",
              "          1.1635e-01,  1.4905e-01, -2.8498e-08,  1.8292e-01, -3.0661e-01,\n",
              "          8.0211e-01,  2.4950e-01, -9.3676e-02,  2.7253e-01,  1.6981e-01,\n",
              "         -3.4504e-01,  1.8387e-01, -4.7821e-08,  1.1757e-01,  1.7852e-01,\n",
              "          8.6810e-02,  1.2619e-01,  8.7449e-01,  1.7493e-01,  2.1305e-01,\n",
              "          2.2845e-01,  1.3461e-01, -2.5138e-08,  2.3123e-01,  1.4359e-01,\n",
              "          2.2817e-01, -7.9652e-08, -4.5994e-08, -3.0792e-08,  1.8274e-01,\n",
              "          1.4877e-01, -5.0358e-03,  3.0524e-01, -3.4730e-08, -1.9635e-01,\n",
              "         -8.1623e-08, -3.8756e-01,  2.4019e-01, -4.7499e-08,  1.8856e-01,\n",
              "         -5.0668e-08,  5.7636e-01,  2.3480e-01,  1.2162e-01,  1.8764e-01,\n",
              "          6.6364e-01,  2.4253e-01, -2.2361e-08,  8.8926e-01]),\n",
              " 'bn1.running_mean': tensor([ 6.7434e-04,  1.3178e-03, -6.6709e-05, -2.4916e-02,  4.2547e-02,\n",
              "          3.8782e-08, -1.6962e-03,  4.0758e-08, -3.2987e-08,  3.1917e-03,\n",
              "         -2.2578e-03, -8.7286e-02,  4.6010e-03,  1.3176e-02, -1.6492e-02,\n",
              "         -1.2173e-03,  1.7546e-03, -3.0290e-08,  1.6268e-04,  7.8634e-04,\n",
              "         -8.9795e-02, -2.6310e-03, -5.1736e-02,  6.3254e-04,  1.3698e-03,\n",
              "          3.2370e-03,  7.3049e-04, -3.3529e-08, -2.2921e-03,  1.5378e-03,\n",
              "          2.3828e-03,  8.1046e-03,  5.9334e-02, -1.1453e-03,  9.6912e-04,\n",
              "          8.1441e-03, -1.0103e-02, -1.6600e-08,  8.4020e-03,  9.1157e-02,\n",
              "         -5.1939e-05, -9.4739e-09,  4.7389e-08, -9.9747e-08, -3.6964e-05,\n",
              "          4.1546e-04, -5.2674e-04, -2.9251e-03,  2.8907e-10, -9.0551e-03,\n",
              "          1.7262e-08, -3.0170e-02,  2.9440e-03,  4.4147e-10, -7.2934e-04,\n",
              "          1.4202e-08,  4.8740e-02,  1.9426e-04,  9.6779e-03, -2.1996e-03,\n",
              "          5.7575e-02, -3.3371e-03,  3.4562e-08, -3.1633e-02]),\n",
              " 'bn1.running_var': tensor([9.7497e-01, 2.3345e-01, 1.7732e+00, 1.7343e+01, 6.0897e+00, 4.8064e-11,\n",
              "         5.8729e-01, 8.6080e-12, 2.3825e-11, 6.3473e-02, 1.9359e+00, 1.1410e+01,\n",
              "         5.5104e-01, 3.6822e+00, 1.1111e+00, 6.8941e+00, 6.8017e+00, 3.3714e-10,\n",
              "         7.7860e-01, 4.0977e+00, 1.4606e+01, 1.3792e+00, 6.7013e+00, 1.4881e-01,\n",
              "         8.6179e-01, 2.2529e+00, 4.7792e-01, 1.4830e-09, 2.0068e+00, 1.4325e+00,\n",
              "         2.0490e+00, 3.4460e+00, 9.9144e+00, 1.4721e+00, 5.3538e-01, 2.3127e-01,\n",
              "         7.4398e+00, 7.1005e-11, 2.7270e-01, 7.7587e+00, 4.5018e-01, 2.9869e-10,\n",
              "         2.0870e-11, 7.9741e-10, 4.4010e-01, 1.4501e+00, 3.7487e-03, 7.3581e-01,\n",
              "         1.4950e-12, 3.2586e+00, 3.6536e-12, 4.1061e+00, 2.6625e+00, 3.0072e-12,\n",
              "         7.6038e-01, 2.4465e-11, 6.7841e+00, 2.9966e-01, 5.8030e+00, 9.4973e-01,\n",
              "         1.4224e+01, 6.5464e-01, 8.5790e-12, 1.8382e+01]),\n",
              " 'bn1.num_batches_tracked': tensor(0)}"
            ]
          },
          "execution_count": 18,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "bn.state_dict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vk4RaHrNx0RN",
        "outputId": "dc41c8c7-98e6-455a-ab8f-888f9505a7d5"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "4"
            ]
          },
          "execution_count": 19,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "len(bn.layers)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3qkqQ3gxx6bL"
      },
      "outputs": [],
      "source": [
        "selected_layers = list(range(2, 8))\n",
        "num_layers = max(selected_layers) + 1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bpXPx2m0yEq1"
      },
      "outputs": [],
      "source": [
        "class FPN(nn.Module):\n",
        "    \"\"\"\n",
        "    https://medium.com/@freshtechyy/fusing-backbone-features-using-feature-pyramid-network-fpn-c652aa6a264b\n",
        "    https://github.com/freshtechyy/resnet/blob/main/resnet_fpn.py\n",
        "    \"\"\"\n",
        "    def __init__(self, in_channels):\n",
        "        super().__init__()\n",
        "        self.interpolation_mode = \"bilinear\"\n",
        "        self.num_features = 256\n",
        "        self.num_downsample = 1\n",
        "        self.pad = True\n",
        "        self.use_conv_downsample = False\n",
        "        self.relu_downsample_layers = False\n",
        "        self.relu_pred_layers = True\n",
        "\n",
        "        self.lat_layers = nn.ModuleList([\n",
        "            nn.Conv2d(x, self.num_features, kernel_size=1)\n",
        "            for x in reversed(in_channels)\n",
        "        ])\n",
        "\n",
        "        padding = 1 if self.pad else 0\n",
        "        self.pred_layers = nn.ModuleList([\n",
        "            nn.Conv2d(self.num_features, self.num_features, kernel_size=3, padding=padding)\n",
        "            for _ in in_channels\n",
        "        ])\n",
        "\n",
        "        if self.use_conv_downsample:\n",
        "            self.downsample_layers = nn.ModuleList([\n",
        "                nn.Conv2d(self.num_features, self.num_features, kernel_size=3, padding=1, stride=2)\n",
        "                for _ in range(self.num_downsample)\n",
        "            ])\n",
        "\n",
        "    def forward(self, convouts):\n",
        "        out = []\n",
        "        x = torch.zeros(1, device=convouts[0].device)\n",
        "\n",
        "        for i, convout in enumerate(reversed(convouts)):\n",
        "            if i > 0:\n",
        "                _, _, h, w = convout.size()\n",
        "                x = F.interpolate(x, size=(h, w), mode=self.interpolation_mode, align_corners=False)\n",
        "            x = x + self.lat_layers[i](convout)\n",
        "            out.append(self.pred_layers[i](x))\n",
        "\n",
        "        if self.relu_pred_layers:\n",
        "            for i in range(len(out)):\n",
        "                out[i] = F.relu(out[i], inplace=True)\n",
        "\n",
        "        if self.use_conv_downsample:\n",
        "            for downsample_layer in self.downsample_layers:\n",
        "                out.append(downsample_layer(out[-1]))\n",
        "        else:\n",
        "            for _ in range(self.num_downsample):\n",
        "                out.append(nn.functional.max_pool2d(out[-1], 1, stride=2))\n",
        "\n",
        "        if self.relu_downsample_layers:\n",
        "            for i in range(-self.num_downsample, 0):\n",
        "                out[i] = F.relu(out[i], inplace=False)\n",
        "\n",
        "        return out[::-1]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oya7cKPL0JQk",
        "outputId": "bf20f1b0-3951-401c-d987-7bf4362e6ee0"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "FPN(\n",
              "  (lat_layers): ModuleList(\n",
              "    (0): Conv2d(2048, 256, kernel_size=(1, 1), stride=(1, 1))\n",
              "    (1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))\n",
              "    (2): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n",
              "    (3): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))\n",
              "  )\n",
              "  (pred_layers): ModuleList(\n",
              "    (0-3): 4 x Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "  )\n",
              ")"
            ]
          },
          "execution_count": 22,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "fpn = FPN([256, 512, 1024, 2048])\n",
        "\n",
        "fpn"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zzuiQ8VS5-qa"
      },
      "outputs": [],
      "source": [
        "class ProtoNet(nn.Module):\n",
        "\n",
        "  def __init__(self):\n",
        "      super().__init__()\n",
        "\n",
        "      self.layer1 = nn.ModuleList()\n",
        "      for i in range(3):\n",
        "            self.layer1.append(nn.Conv2d(256, 256, kernel_size=3, padding=1))\n",
        "            self.layer1.append(nn.ReLU(inplace=True))\n",
        "      self.layer2 = nn.Conv2d(256, 256, kernel_size=3, padding=1)\n",
        "      self.relu = nn.ReLU(inplace=True)\n",
        "      self.layer3 = nn.Conv2d(256, 32, kernel_size=1)\n",
        "\n",
        "  def forward(self,x):\n",
        "      out = x\n",
        "      for layer in self.layer1:\n",
        "        out=layer(out)\n",
        "      out = F.interpolate(out, (138, 138), mode='bilinear', align_corners=False)\n",
        "      out = self.relu(out)\n",
        "      out = self.layer2(out)\n",
        "      out = self.relu(out)\n",
        "      out = self.layer3(out)\n",
        "      out = self.relu(out)\n",
        "\n",
        "      return out"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QC7f3JEI5_Uh"
      },
      "outputs": [],
      "source": [
        "class PredictionNet(nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "\n",
        "    def forward(self,x):\n",
        "        pass"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0cslSOeX3KLY"
      },
      "outputs": [],
      "source": [
        "class Yolact(torch.nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        # setup backbone\n",
        "        self.backbone = ResNetBackbone([3, 4, 23, 3]) # resnet101\n",
        "\n",
        "        self.selected_layers = list(range(2, 8))\n",
        "        num_layers = max(selected_layers) + 1\n",
        "\n",
        "        while len(self.backbone.layers) < num_layers:\n",
        "            self.backbone.add_layer()\n",
        "\n",
        "        # freeze backbone\n",
        "\n",
        "        # init fpn\n",
        "        src_channels = self.backbone.channels\n",
        "        self.fpn = FPN([src_channels[i] for i in self.selected_layers])\n",
        "        self.selected_layers = list(range(len(self.selected_layers) + self.fpn.num_downsample))\n",
        "        src_channels = [self.fpn.num_features] * len(self.selected_layers)\n",
        "\n",
        "        # retina_net = retinanet_resnet50_fpn(weights=RetinaNet_ResNet50_FPN_Weights)\n",
        "        # self.retina_net_backbone = retina_net.backbone\n",
        "\n",
        "        # init protonet\n",
        "        self.proto_net = ProtoNet()\n",
        "\n",
        "        # init prediction net\n",
        "        self.pred_net = PredictionNet()\n",
        "\n",
        "    def forward(self, x):\n",
        "        outs = self.backbone(x)\n",
        "        outs = [outs[i] for i in self.selected_layers]\n",
        "\n",
        "        # outs = self.retina_net_backbone(outs)\n",
        "        # proto_out = self.proto_net(outs[0])\n",
        "\n",
        "        outs = self.fpn(outs)\n",
        "\n",
        "        # proto_out = self.proto_net(outs[0])\n",
        "\n",
        "        # return proto_out\n",
        "\n",
        "        return outs\n",
        "\n",
        "    def init_backbone_weights(self):  # TODO: make sure names map one to one\n",
        "        resnet101 = torchvision.models.resnet101(\n",
        "            weights=torchvision.models.resnet101_weights.ResNet101_Weights.DEFAULT\n",
        "        )\n",
        "        state_dict = resnet101.state_dict()\n",
        "\n",
        "        # TODO: strict=True?\n",
        "        self.backbone.load_state_dict(state_dict, strict=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "N6ZJG7KFyS1I",
        "outputId": "534d1a79-fcef-4140-d404-92c2e7eb84ff"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[256, 512, 1024, 2048, 1024, 1024, 1024, 1024]"
            ]
          },
          "execution_count": 26,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "yolact = Yolact()\n",
        "len(yolact.backbone.layers)\n",
        "yolact.backbone.channels"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_q3TSOyd6sVW",
        "outputId": "dffa79e8-a4a7-4877-c45a-3ae5931ffa98"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "FPN(\n",
              "  (lat_layers): ModuleList(\n",
              "    (0-3): 4 x Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))\n",
              "    (4): Conv2d(2048, 256, kernel_size=(1, 1), stride=(1, 1))\n",
              "    (5): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))\n",
              "  )\n",
              "  (pred_layers): ModuleList(\n",
              "    (0-5): 6 x Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "  )\n",
              ")"
            ]
          },
          "execution_count": 27,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "fpn = FPN([yolact.backbone.channels[i] for i in selected_layers])\n",
        "\n",
        "fpn"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 347
        },
        "id": "5pRzjk9F74UH",
        "outputId": "54c1c98e-f6b1-445d-bd99-2f768ea3679a"
      },
      "outputs": [
        {
          "ename": "RuntimeError",
          "evalue": "Given groups=1, weight of size [256, 1024, 1, 1], expected input[1, 2048, 16, 25] to have 1024 channels, but got 2048 channels instead",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-28-e64ff7ccbb20>\u001b[0m in \u001b[0;36m<cell line: 3>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0myolact\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meval\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0mt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m512\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m800\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m \u001b[0myolact\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mt\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      4\u001b[0m \u001b[0;31m# len(yolact.backbone(t))\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1509\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1510\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1511\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\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   1512\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1513\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1518\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1519\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1520\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\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   1521\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1522\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-25-a1d2db99aa3c>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     36\u001b[0m         \u001b[0;31m# proto_out = self.proto_net(outs[0])\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     37\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 38\u001b[0;31m         \u001b[0mouts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfpn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mouts\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     39\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     40\u001b[0m         \u001b[0;31m# proto_out = self.proto_net(outs[0])\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1509\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1510\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1511\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\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   1512\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1513\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1518\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1519\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1520\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\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   1521\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1522\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-21-7d224d589ed6>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, convouts)\u001b[0m\n\u001b[1;32m     39\u001b[0m                 \u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mh\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mw\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconvout\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     40\u001b[0m                 \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minterpolate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msize\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mh\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mw\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmode\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minterpolation_mode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0malign_corners\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 41\u001b[0;31m             \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mx\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlat_layers\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mconvout\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     42\u001b[0m             \u001b[0mout\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpred_layers\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     43\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1509\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1510\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1511\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\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   1512\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1513\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1518\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1519\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1520\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\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   1521\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1522\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/conv.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    458\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    459\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 460\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_conv_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mweight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbias\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    461\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    462\u001b[0m \u001b[0;32mclass\u001b[0m \u001b[0mConv3d\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_ConvNd\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/conv.py\u001b[0m in \u001b[0;36m_conv_forward\u001b[0;34m(self, input, weight, bias)\u001b[0m\n\u001b[1;32m    454\u001b[0m                             \u001b[0mweight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbias\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstride\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    455\u001b[0m                             _pair(0), self.dilation, self.groups)\n\u001b[0;32m--> 456\u001b[0;31m         return F.conv2d(input, weight, bias, self.stride,\n\u001b[0m\u001b[1;32m    457\u001b[0m                         self.padding, self.dilation, self.groups)\n\u001b[1;32m    458\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mRuntimeError\u001b[0m: Given groups=1, weight of size [256, 1024, 1, 1], expected input[1, 2048, 16, 25] to have 1024 channels, but got 2048 channels instead"
          ]
        }
      ],
      "source": [
        "yolact.eval()\n",
        "t = torch.randn(1, 3, 512, 800)\n",
        "yolact(t)\n",
        "# len(yolact.backbone(t))"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Training objective\n",
        "\n",
        "#### Multibox Loss\n",
        "Weighted sum of localization loss and confidence loss"
      ],
      "metadata": {
        "id": "-x4F11g7F90S"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "activation_func = {\n",
        "    'tanh':    torch.tanh,\n",
        "    'sigmoid': torch.sigmoid,\n",
        "    'softmax': lambda x: torch.nn.functional.softmax(x, dim=-1),\n",
        "    'relu':    lambda x: torch.nn.functional.relu(x, inplace=True),\n",
        "    'none':    lambda x: x,\n",
        "}"
      ],
      "metadata": {
        "id": "95a2LtXuceGa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "mask_type = {\n",
        "    \"direct\": 0,\n",
        "    \"lincomb\": 1,\n",
        "}"
      ],
      "metadata": {
        "id": "nwXWIweUcoUw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "\n",
        "@torch.jit.script\n",
        "def point_form(boxes):\n",
        "    \"\"\" Convert prior_boxes to (xmin, ymin, xmax, ymax)\n",
        "    representation for comparison to point form ground truth data.\n",
        "    Args:\n",
        "        boxes: (tensor) center-size default boxes from priorbox layers.\n",
        "    Return:\n",
        "        boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.\n",
        "    \"\"\"\n",
        "    return torch.cat((boxes[:, :2] - boxes[:, 2:]/2,     # xmin, ymin\n",
        "                     boxes[:, :2] + boxes[:, 2:]/2), 1)  # xmax, ymax\n",
        "\n",
        "\n",
        "@torch.jit.script\n",
        "def center_size(boxes):\n",
        "    \"\"\" Convert prior_boxes to (cx, cy, w, h)\n",
        "    representation for comparison to center-size form ground truth data.\n",
        "    Args:\n",
        "        boxes: (tensor) point_form boxes\n",
        "    Return:\n",
        "        boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.\n",
        "    \"\"\"\n",
        "    return torch.cat(( (boxes[:, 2:] + boxes[:, :2])/2,     # cx, cy\n",
        "                        boxes[:, 2:] - boxes[:, :2]  ), 1)  # w, h\n",
        "\n",
        "@torch.jit.script\n",
        "def intersect(box_a, box_b):\n",
        "    \"\"\" We resize both tensors to [A,B,2] without new malloc:\n",
        "    [A,2] -> [A,1,2] -> [A,B,2]\n",
        "    [B,2] -> [1,B,2] -> [A,B,2]\n",
        "    Then we compute the area of intersect between box_a and box_b.\n",
        "    Args:\n",
        "      box_a: (tensor) bounding boxes, Shape: [n,A,4].\n",
        "      box_b: (tensor) bounding boxes, Shape: [n,B,4].\n",
        "    Return:\n",
        "      (tensor) intersection area, Shape: [n,A,B].\n",
        "    \"\"\"\n",
        "    n = box_a.size(0)\n",
        "    A = box_a.size(1)\n",
        "    B = box_b.size(1)\n",
        "    max_xy = torch.min(box_a[:, :, 2:].unsqueeze(2).expand(n, A, B, 2),\n",
        "                       box_b[:, :, 2:].unsqueeze(1).expand(n, A, B, 2))\n",
        "    min_xy = torch.max(box_a[:, :, :2].unsqueeze(2).expand(n, A, B, 2),\n",
        "                       box_b[:, :, :2].unsqueeze(1).expand(n, A, B, 2))\n",
        "    return torch.clamp(max_xy - min_xy, min=0).prod(3)  # inter\n",
        "\n",
        "\n",
        "def jaccard(box_a, box_b, iscrowd:bool=False):\n",
        "    \"\"\"Compute the jaccard overlap of two sets of boxes.  The jaccard overlap\n",
        "    is simply the intersection over union of two boxes.  Here we operate on\n",
        "    ground truth boxes and default boxes. If iscrowd=True, put the crowd in box_b.\n",
        "    E.g.:\n",
        "        A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)\n",
        "    Args:\n",
        "        box_a: (tensor) Ground truth bounding boxes, Shape: [num_objects,4]\n",
        "        box_b: (tensor) Prior boxes from priorbox layers, Shape: [num_priors,4]\n",
        "    Return:\n",
        "        jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)]\n",
        "    \"\"\"\n",
        "    use_batch = True\n",
        "    if box_a.dim() == 2:\n",
        "        use_batch = False\n",
        "        box_a = box_a[None, ...]\n",
        "        box_b = box_b[None, ...]\n",
        "\n",
        "    inter = intersect(box_a, box_b)\n",
        "    area_a = ((box_a[:, :, 2]-box_a[:, :, 0]) *\n",
        "              (box_a[:, :, 3]-box_a[:, :, 1])).unsqueeze(2).expand_as(inter)  # [A,B]\n",
        "    area_b = ((box_b[:, :, 2]-box_b[:, :, 0]) *\n",
        "              (box_b[:, :, 3]-box_b[:, :, 1])).unsqueeze(1).expand_as(inter)  # [A,B]\n",
        "    union = area_a + area_b - inter\n",
        "\n",
        "    out = inter / area_a if iscrowd else inter / union\n",
        "    return out if use_batch else out.squeeze(0)\n",
        "\n",
        "def elemwise_box_iou(box_a, box_b):\n",
        "    \"\"\" Does the same as above but instead of pairwise, elementwise along the inner dimension. \"\"\"\n",
        "    max_xy = torch.min(box_a[:, 2:], box_b[:, 2:])\n",
        "    min_xy = torch.max(box_a[:, :2], box_b[:, :2])\n",
        "    inter = torch.clamp((max_xy - min_xy), min=0)\n",
        "    inter = inter[:, 0] * inter[:, 1]\n",
        "\n",
        "    area_a = (box_a[:, 2] - box_a[:, 0]) * (box_a[:, 3] - box_a[:, 1])\n",
        "    area_b = (box_b[:, 2] - box_b[:, 0]) * (box_b[:, 3] - box_b[:, 1])\n",
        "\n",
        "    union = area_a + area_b - inter\n",
        "    union = torch.clamp(union, min=0.1)\n",
        "\n",
        "    # Return value is [n] for inputs [n, 4]\n",
        "    return torch.clamp(inter / union, max=1)\n",
        "\n",
        "def mask_iou(masks_a, masks_b, iscrowd=False):\n",
        "    \"\"\"\n",
        "    Computes the pariwise mask IoU between two sets of masks of size [a, h, w] and [b, h, w].\n",
        "    The output is of size [a, b].\n",
        "\n",
        "    Wait I thought this was \"box_utils\", why am I putting this in here?\n",
        "    \"\"\"\n",
        "\n",
        "    masks_a = masks_a.view(masks_a.size(0), -1)\n",
        "    masks_b = masks_b.view(masks_b.size(0), -1)\n",
        "\n",
        "    intersection = masks_a @ masks_b.t()\n",
        "    area_a = masks_a.sum(dim=1).unsqueeze(1)\n",
        "    area_b = masks_b.sum(dim=1).unsqueeze(0)\n",
        "\n",
        "    return intersection / (area_a + area_b - intersection) if not iscrowd else intersection / area_a\n",
        "\n",
        "def elemwise_mask_iou(masks_a, masks_b):\n",
        "    \"\"\" Does the same as above but instead of pairwise, elementwise along the outer dimension. \"\"\"\n",
        "    masks_a = masks_a.view(-1, masks_a.size(-1))\n",
        "    masks_b = masks_b.view(-1, masks_b.size(-1))\n",
        "\n",
        "    intersection = (masks_a * masks_b).sum(dim=0)\n",
        "    area_a = masks_a.sum(dim=0)\n",
        "    area_b = masks_b.sum(dim=0)\n",
        "\n",
        "    # Return value is [n] for inputs [h, w, n]\n",
        "    return torch.clamp(intersection / torch.clamp(area_a + area_b - intersection, min=0.1), max=1)\n",
        "\n",
        "\n",
        "\n",
        "def change(gt, priors):\n",
        "    \"\"\"\n",
        "    Compute the d_change metric proposed in Box2Pix:\n",
        "    https://lmb.informatik.uni-freiburg.de/Publications/2018/UB18/paper-box2pix.pdf\n",
        "\n",
        "    Input should be in point form (xmin, ymin, xmax, ymax).\n",
        "\n",
        "    Output is of shape [num_gt, num_priors]\n",
        "    Note this returns -change so it can be a drop in replacement for\n",
        "    \"\"\"\n",
        "    num_priors = priors.size(0)\n",
        "    num_gt     = gt.size(0)\n",
        "\n",
        "    gt_w = (gt[:, 2] - gt[:, 0])[:, None].expand(num_gt, num_priors)\n",
        "    gt_h = (gt[:, 3] - gt[:, 1])[:, None].expand(num_gt, num_priors)\n",
        "\n",
        "    gt_mat =     gt[:, None, :].expand(num_gt, num_priors, 4)\n",
        "    pr_mat = priors[None, :, :].expand(num_gt, num_priors, 4)\n",
        "\n",
        "    diff = gt_mat - pr_mat\n",
        "    diff[:, :, 0] /= gt_w\n",
        "    diff[:, :, 2] /= gt_w\n",
        "    diff[:, :, 1] /= gt_h\n",
        "    diff[:, :, 3] /= gt_h\n",
        "\n",
        "    return -torch.sqrt( (diff ** 2).sum(dim=2) )\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "def match(pos_thresh, neg_thresh, truths, priors, labels, crowd_boxes, loc_t, conf_t, idx_t, idx, loc_data):\n",
        "    \"\"\"Match each prior box with the ground truth box of the highest jaccard\n",
        "    overlap, encode the bounding boxes, then return the matched indices\n",
        "    corresponding to both confidence and location preds.\n",
        "    Args:\n",
        "        pos_thresh: (float) IoU > pos_thresh ==> positive.\n",
        "        neg_thresh: (float) IoU < neg_thresh ==> negative.\n",
        "        truths: (tensor) Ground truth boxes, Shape: [num_obj, num_priors].\n",
        "        priors: (tensor) Prior boxes from priorbox layers, Shape: [n_priors,4].\n",
        "        labels: (tensor) All the class labels for the image, Shape: [num_obj].\n",
        "        crowd_boxes: (tensor) All the crowd box annotations or None if there are none.\n",
        "        loc_t: (tensor) Tensor to be filled w/ endcoded location targets.\n",
        "        conf_t: (tensor) Tensor to be filled w/ matched indices for conf preds. Note: -1 means neutral.\n",
        "        idx_t: (tensor) Tensor to be filled w/ the index of the matched gt box for each prior.\n",
        "        idx: (int) current batch index.\n",
        "        loc_data: (tensor) The predicted bbox regression coordinates for this batch.\n",
        "    Return:\n",
        "        The matched indices corresponding to 1)location and 2)confidence preds.\n",
        "    \"\"\"\n",
        "    use_yolo_regressors = True\n",
        "    use_prediction_matching = True\n",
        "    use_change_matching = False\n",
        "    crowd_iou_threshold = 0.5\n",
        "\n",
        "\n",
        "    decoded_priors = decode(loc_data, priors, use_yolo_regressors) if use_prediction_matching else point_form(priors)\n",
        "\n",
        "    # Size [num_objects, num_priors]\n",
        "    overlaps = jaccard(truths, decoded_priors) if not use_change_matching else change(truths, decoded_priors)\n",
        "\n",
        "    # Size [num_priors] best ground truth for each prior\n",
        "    best_truth_overlap, best_truth_idx = overlaps.max(0)\n",
        "\n",
        "    # We want to ensure that each gt gets used at least once so that we don't\n",
        "    # waste any training data. In order to do that, find the max overlap anchor\n",
        "    # with each gt, and force that anchor to use that gt.\n",
        "    for _ in range(overlaps.size(0)):\n",
        "        # Find j, the gt with the highest overlap with a prior\n",
        "        # In effect, this will loop through overlaps.size(0) in a \"smart\" order,\n",
        "        # always choosing the highest overlap first.\n",
        "        best_prior_overlap, best_prior_idx = overlaps.max(1)\n",
        "        j = best_prior_overlap.max(0)[1]\n",
        "\n",
        "        # Find i, the highest overlap anchor with this gt\n",
        "        i = best_prior_idx[j]\n",
        "\n",
        "        # Set all other overlaps with i to be -1 so that no other gt uses it\n",
        "        overlaps[:, i] = -1\n",
        "        # Set all other overlaps with j to be -1 so that this loop never uses j again\n",
        "        overlaps[j, :] = -1\n",
        "\n",
        "        # Overwrite i's score to be 2 so it doesn't get thresholded ever\n",
        "        best_truth_overlap[i] = 2\n",
        "        # Set the gt to be used for i to be j, overwriting whatever was there\n",
        "        best_truth_idx[i] = j\n",
        "\n",
        "    matches = truths[best_truth_idx]            # Shape: [num_priors,4]\n",
        "    conf = labels[best_truth_idx] + 1           # Shape: [num_priors]\n",
        "\n",
        "    conf[best_truth_overlap < pos_thresh] = -1  # label as neutral\n",
        "    conf[best_truth_overlap < neg_thresh] =  0  # label as background\n",
        "\n",
        "    # Deal with crowd annotations for COCO\n",
        "    if crowd_boxes is not None and crowd_iou_threshold < 1:\n",
        "        # Size [num_priors, num_crowds]\n",
        "        crowd_overlaps = jaccard(decoded_priors, crowd_boxes, iscrowd=True)\n",
        "        # Size [num_priors]\n",
        "        best_crowd_overlap, best_crowd_idx = crowd_overlaps.max(1)\n",
        "        # Set non-positives with crowd iou of over the threshold to be neutral.\n",
        "        conf[(conf <= 0) & (best_crowd_overlap > crowd_iou_threshold)] = -1\n",
        "\n",
        "    loc = encode(matches, priors, use_yolo_regressors)\n",
        "    loc_t[idx]  = loc    # [num_priors,4] encoded offsets to learn\n",
        "    conf_t[idx] = conf   # [num_priors] top class label for each prior\n",
        "    idx_t[idx]  = best_truth_idx # [num_priors] indices for lookup\n",
        "\n",
        "@torch.jit.script\n",
        "def encode(matched, priors, use_yolo_regressors:bool=False):\n",
        "    \"\"\"\n",
        "    Encode bboxes matched with each prior into the format\n",
        "    produced by the network. See decode for more details on\n",
        "    this format. Note that encode(decode(x, p), p) = x.\n",
        "\n",
        "    Args:\n",
        "        - matched: A tensor of bboxes in point form with shape [num_priors, 4]\n",
        "        - priors:  The tensor of all priors with shape [num_priors, 4]\n",
        "    Return: A tensor with encoded relative coordinates in the format\n",
        "            outputted by the network (see decode). Size: [num_priors, 4]\n",
        "    \"\"\"\n",
        "\n",
        "    if use_yolo_regressors:\n",
        "        # Exactly the reverse of what we did in decode\n",
        "        # In fact encode(decode(x, p), p) should be x\n",
        "        boxes = center_size(matched)\n",
        "\n",
        "        loc = torch.cat((\n",
        "            boxes[:, :2] - priors[:, :2],\n",
        "            torch.log(boxes[:, 2:] / priors[:, 2:])\n",
        "        ), 1)\n",
        "    else:\n",
        "        variances = [0.1, 0.2]\n",
        "\n",
        "        # dist b/t match center and prior's center\n",
        "        g_cxcy = (matched[:, :2] + matched[:, 2:])/2 - priors[:, :2]\n",
        "        # encode variance\n",
        "        g_cxcy /= (variances[0] * priors[:, 2:])\n",
        "        # match wh / prior wh\n",
        "        g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:]\n",
        "        g_wh = torch.log(g_wh) / variances[1]\n",
        "        # return target for smooth_l1_loss\n",
        "        loc = torch.cat([g_cxcy, g_wh], 1)  # [num_priors,4]\n",
        "\n",
        "    return loc\n",
        "\n",
        "@torch.jit.script\n",
        "def decode(loc, priors, use_yolo_regressors:bool=False):\n",
        "    \"\"\"\n",
        "    Decode predicted bbox coordinates using the same scheme\n",
        "    employed by Yolov2: https://arxiv.org/pdf/1612.08242.pdf\n",
        "\n",
        "        b_x = (sigmoid(pred_x) - .5) / conv_w + prior_x\n",
        "        b_y = (sigmoid(pred_y) - .5) / conv_h + prior_y\n",
        "        b_w = prior_w * exp(loc_w)\n",
        "        b_h = prior_h * exp(loc_h)\n",
        "\n",
        "    Note that loc is inputed as [(s(x)-.5)/conv_w, (s(y)-.5)/conv_h, w, h]\n",
        "    while priors are inputed as [x, y, w, h] where each coordinate\n",
        "    is relative to size of the image (even sigmoid(x)). We do this\n",
        "    in the network by dividing by the 'cell size', which is just\n",
        "    the size of the convouts.\n",
        "\n",
        "    Also note that prior_x and prior_y are center coordinates which\n",
        "    is why we have to subtract .5 from sigmoid(pred_x and pred_y).\n",
        "\n",
        "    Args:\n",
        "        - loc:    The predicted bounding boxes of size [num_priors, 4]\n",
        "        - priors: The priorbox coords with size [num_priors, 4]\n",
        "\n",
        "    Returns: A tensor of decoded relative coordinates in point form\n",
        "             form with size [num_priors, 4]\n",
        "    \"\"\"\n",
        "\n",
        "    if use_yolo_regressors:\n",
        "        # Decoded boxes in center-size notation\n",
        "        boxes = torch.cat((\n",
        "            loc[:, :2] + priors[:, :2],\n",
        "            priors[:, 2:] * torch.exp(loc[:, 2:])\n",
        "        ), 1)\n",
        "\n",
        "        boxes = point_form(boxes)\n",
        "    else:\n",
        "        variances = [0.1, 0.2]\n",
        "\n",
        "        boxes = torch.cat((\n",
        "            priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],\n",
        "            priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)\n",
        "        boxes[:, :2] -= boxes[:, 2:] / 2\n",
        "        boxes[:, 2:] += boxes[:, :2]\n",
        "\n",
        "    return boxes\n",
        "\n",
        "\n",
        "\n",
        "def log_sum_exp(x):\n",
        "    \"\"\"Utility function for computing log_sum_exp while determining\n",
        "    This will be used to determine unaveraged confidence loss across\n",
        "    all examples in a batch.\n",
        "    Args:\n",
        "        x (Variable(tensor)): conf_preds from conf layers\n",
        "    \"\"\"\n",
        "    x_max = x.data.max()\n",
        "    return torch.log(torch.sum(torch.exp(x-x_max), 1)) + x_max\n",
        "\n",
        "\n",
        "@torch.jit.script\n",
        "def sanitize_coordinates(_x1, _x2, img_size:int, padding:int=0, cast:bool=True):\n",
        "    \"\"\"\n",
        "    Sanitizes the input coordinates so that x1 < x2, x1 != x2, x1 >= 0, and x2 <= image_size.\n",
        "    Also converts from relative to absolute coordinates and casts the results to long tensors.\n",
        "\n",
        "    If cast is false, the result won't be cast to longs.\n",
        "    Warning: this does things in-place behind the scenes so copy if necessary.\n",
        "    \"\"\"\n",
        "    _x1 = _x1 * img_size\n",
        "    _x2 = _x2 * img_size\n",
        "    if cast:\n",
        "        _x1 = _x1.long()\n",
        "        _x2 = _x2.long()\n",
        "    x1 = torch.min(_x1, _x2)\n",
        "    x2 = torch.max(_x1, _x2)\n",
        "    x1 = torch.clamp(x1-padding, min=0)\n",
        "    x2 = torch.clamp(x2+padding, max=img_size)\n",
        "\n",
        "    return x1, x2\n",
        "\n",
        "\n",
        "@torch.jit.script\n",
        "def crop(masks, boxes, padding:int=1):\n",
        "    \"\"\"\n",
        "    \"Crop\" predicted masks by zeroing out everything not in the predicted bbox.\n",
        "    Vectorized by Chong (thanks Chong).\n",
        "\n",
        "    Args:\n",
        "        - masks should be a size [h, w, n] tensor of masks\n",
        "        - boxes should be a size [n, 4] tensor of bbox coords in relative point form\n",
        "    \"\"\"\n",
        "    h, w, n = masks.size()\n",
        "    x1, x2 = sanitize_coordinates(boxes[:, 0], boxes[:, 2], w, padding, cast=False)\n",
        "    y1, y2 = sanitize_coordinates(boxes[:, 1], boxes[:, 3], h, padding, cast=False)\n",
        "\n",
        "    rows = torch.arange(w, device=masks.device, dtype=x1.dtype).view(1, -1, 1).expand(h, w, n)\n",
        "    cols = torch.arange(h, device=masks.device, dtype=x1.dtype).view(-1, 1, 1).expand(h, w, n)\n",
        "\n",
        "    masks_left  = rows >= x1.view(1, 1, -1)\n",
        "    masks_right = rows <  x2.view(1, 1, -1)\n",
        "    masks_up    = cols >= y1.view(1, 1, -1)\n",
        "    masks_down  = cols <  y2.view(1, 1, -1)\n",
        "\n",
        "    crop_mask = masks_left * masks_right * masks_up * masks_down\n",
        "\n",
        "    return masks * crop_mask.float()\n",
        "\n",
        "\n",
        "def index2d(src, idx):\n",
        "    \"\"\"\n",
        "    Indexes a tensor by a 2d index.\n",
        "\n",
        "    In effect, this does\n",
        "        out[i, j] = src[i, idx[i, j]]\n",
        "\n",
        "    Both src and idx should have the same size.\n",
        "    \"\"\"\n",
        "\n",
        "    offs = torch.arange(idx.size(0), device=idx.device)[:, None].expand_as(idx)\n",
        "    idx  = idx + offs * idx.size(1)\n",
        "\n",
        "    return src.view(-1)[idx.view(-1)].view(idx.size())"
      ],
      "metadata": {
        "id": "1c-jBwvZbm20"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# -*- coding: utf-8 -*-\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.autograd import Variable\n",
        "\n",
        "class MultiBoxLoss(nn.Module):\n",
        "    \"\"\"SSD Weighted Loss Function\n",
        "    Compute Targets:\n",
        "        1) Produce Confidence Target Indices by matching  ground truth boxes\n",
        "           with (default) 'priorboxes' that have jaccard index > threshold parameter\n",
        "           (default threshold: 0.5).\n",
        "        2) Produce localization target by 'encoding' variance into offsets of ground\n",
        "           truth boxes and their matched  'priorboxes'.\n",
        "        3) Hard negative mining to filter the excessive number of negative examples\n",
        "           that comes with using a large number of default bounding boxes.\n",
        "           (default negative:positive ratio 3:1)\n",
        "    Objective Loss:\n",
        "        L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N\n",
        "        Where, Lconf is the CrossEntropy Loss and Lloc is the SmoothL1 Loss\n",
        "        weighted by α which is set to 1 by cross val.\n",
        "        Args:\n",
        "            c: class confidences,\n",
        "            l: predicted boxes,\n",
        "            g: ground truth boxes\n",
        "            N: number of matched default boxes\n",
        "        See: https://arxiv.org/pdf/1512.02325.pdf for more details.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, num_classes, pos_threshold, neg_threshold, negpos_ratio):\n",
        "        super(MultiBoxLoss, self).__init__()\n",
        "        self.num_classes = num_classes\n",
        "\n",
        "        self.pos_threshold = pos_threshold\n",
        "        self.neg_threshold = neg_threshold\n",
        "        self.negpos_ratio = negpos_ratio\n",
        "\n",
        "        use_class_balanced_conf = False\n",
        "\n",
        "        self.mask_type = mask_type[\"lincomb\"]\n",
        "        self.use_mask_scoring = False\n",
        "        self.use_instance_coeff = False\n",
        "        self.use_class_existence_loss = False\n",
        "        self.use_class_existence_loss = False\n",
        "        self.train_boxes = True\n",
        "        self.train_masks = True\n",
        "        self.use_gt_bboxes = False\n",
        "        self.use_maskiou = False\n",
        "        self.mask_proto_loss = None\n",
        "        self.use_focal_loss = False\n",
        "        self.mask_dim = None\n",
        "        self.mask_alpha = 6.125\n",
        "        self.bbox_alpha = 1.5\n",
        "        self.use_focal_loss = False\n",
        "        self.use_sigmoid_focal_loss = False\n",
        "        self.use_objectness_score = False\n",
        "        self.maskiou_alpha = 25\n",
        "        self.maskious_to_train = -1\n",
        "        self.masks_to_train = 300\n",
        "        self.mask_proto_coeff_diversity_loss = False\n",
        "        self.discard_mask_area = 5*5\n",
        "        self.mask_proto_crop = True\n",
        "        self.mask_proto_normalize_emulate_roi_pooling = False\n",
        "        self.mask_proto_reweight_mask_loss = True\n",
        "        self.mask_proto_normalize_mask_loss_by_sqrt_area = False\n",
        "        self.mask_proto_mask_activation = activation_func[\"sigmoid\"]\n",
        "        self.mask_proto_double_loss = False\n",
        "        self.mask_proto_crop_with_pred_box = False\n",
        "        self.mask_proto_reweight_coeff = 1\n",
        "        self.use_yolo_regressors = False\n",
        "        self.mask_proto_binarize_downsampled_gt = True\n",
        "        self.mask_proto_coeff_diversity_alpha = 1\n",
        "        self.conf_alpha = 1\n",
        "        self.mask_size = 16\n",
        "        self.focal_loss_gamma = 2\n",
        "        self.focal_loss_alpha = 0.25\n",
        "        self.use_class_balanced_conf = False\n",
        "        self.ohem_use_most_confident = False\n",
        "        self.semantic_segmentation_alpha = 1\n",
        "        self.class_existence_alpha = 1\n",
        "        self.use_class_existence_loss = True\n",
        "        self.use_semantic_segmentation_loss = True\n",
        "        self.use_objectness_score = True\n",
        "        self.use_sigmoid_focal_loss = True\n",
        "        self.use_focal_loss = True\n",
        "\n",
        "        # If you output a proto mask with this area, your l1 loss will be l1_alpha\n",
        "        # Note that the area is relative (so 1 would be the entire image)\n",
        "        self.l1_expected_area = 20*20/70/70\n",
        "        self.l1_alpha = 0.1\n",
        "\n",
        "        if use_class_balanced_conf:\n",
        "            self.class_instances = None\n",
        "            self.total_instances = 0\n",
        "\n",
        "    def forward(self, net, predictions, targets, masks, num_crowds):\n",
        "        \"\"\"Multibox Loss\n",
        "        Args:\n",
        "            predictions (tuple): A tuple containing loc preds, conf preds,\n",
        "            mask preds, and prior boxes from SSD net.\n",
        "                loc shape: torch.size(batch_size,num_priors,4)\n",
        "                conf shape: torch.size(batch_size,num_priors,num_classes)\n",
        "                masks shape: torch.size(batch_size,num_priors,mask_dim)\n",
        "                priors shape: torch.size(num_priors,4)\n",
        "                proto* shape: torch.size(batch_size,mask_h,mask_w,mask_dim)\n",
        "\n",
        "            targets (list<tensor>): Ground truth boxes and labels for a batch,\n",
        "                shape: [batch_size][num_objs,5] (last idx is the label).\n",
        "\n",
        "            masks (list<tensor>): Ground truth masks for each object in each image,\n",
        "                shape: [batch_size][num_objs,im_height,im_width]\n",
        "\n",
        "            num_crowds (list<int>): Number of crowd annotations per batch. The crowd\n",
        "                annotations should be the last num_crowds elements of targets and masks.\n",
        "\n",
        "            * Only if mask_type == lincomb\n",
        "        \"\"\"\n",
        "\n",
        "        loc_data  = predictions['loc']\n",
        "        conf_data = predictions['conf']\n",
        "        mask_data = predictions['mask']\n",
        "        priors    = predictions['priors']\n",
        "\n",
        "        if self.mask_type == mask_type[\"lincomb\"]:\n",
        "            proto_data = predictions['proto']\n",
        "\n",
        "        score_data = predictions['score'] if self.use_mask_scoring   else None\n",
        "        inst_data  = predictions['inst']  if self.use_instance_coeff else None\n",
        "\n",
        "        labels = [None] * len(targets) # Used in sem segm loss\n",
        "\n",
        "        batch_size = loc_data.size(0)\n",
        "        num_priors = priors.size(0)\n",
        "        num_classes = self.num_classes\n",
        "\n",
        "        # Match priors (default boxes) and ground truth boxes\n",
        "        # These tensors will be created with the same device as loc_data\n",
        "        loc_t = loc_data.new(batch_size, num_priors, 4)\n",
        "        gt_box_t = loc_data.new(batch_size, num_priors, 4)\n",
        "        conf_t = loc_data.new(batch_size, num_priors).long()\n",
        "        idx_t = loc_data.new(batch_size, num_priors).long()\n",
        "\n",
        "        if self.use_class_existence_loss:\n",
        "            class_existence_t = loc_data.new(batch_size, num_classes-1)\n",
        "\n",
        "        for idx in range(batch_size):\n",
        "            truths      = targets[idx][:, :-1].data\n",
        "            labels[idx] = targets[idx][:, -1].data.long()\n",
        "\n",
        "            if self.use_class_existence_loss:\n",
        "                # Construct a one-hot vector for each object and collapse it into an existence vector with max\n",
        "                # Also it's fine to include the crowd annotations here\n",
        "                class_existence_t[idx, :] = torch.eye(num_classes-1, device=conf_t.get_device())[labels[idx]].max(dim=0)[0]\n",
        "\n",
        "            # Split the crowd annotations because they come bundled in\n",
        "            cur_crowds = num_crowds[idx]\n",
        "            if cur_crowds > 0:\n",
        "                split = lambda x: (x[-cur_crowds:], x[:-cur_crowds])\n",
        "                crowd_boxes, truths = split(truths)\n",
        "\n",
        "                # We don't use the crowd labels or masks\n",
        "                _, labels[idx] = split(labels[idx])\n",
        "                _, masks[idx]  = split(masks[idx])\n",
        "            else:\n",
        "                crowd_boxes = None\n",
        "\n",
        "\n",
        "            match(self.pos_threshold, self.neg_threshold,\n",
        "                  truths, priors.data, labels[idx], crowd_boxes,\n",
        "                  loc_t, conf_t, idx_t, idx, loc_data[idx])\n",
        "\n",
        "            gt_box_t[idx, :, :] = truths[idx_t[idx]]\n",
        "\n",
        "        # wrap targets\n",
        "        loc_t = Variable(loc_t, requires_grad=False)\n",
        "        conf_t = Variable(conf_t, requires_grad=False)\n",
        "        idx_t = Variable(idx_t, requires_grad=False)\n",
        "\n",
        "        pos = conf_t > 0\n",
        "        num_pos = pos.sum(dim=1, keepdim=True)\n",
        "\n",
        "        # Shape: [batch,num_priors,4]\n",
        "        pos_idx = pos.unsqueeze(pos.dim()).expand_as(loc_data)\n",
        "\n",
        "        losses = {}\n",
        "\n",
        "        # Localization Loss (Smooth L1)\n",
        "        if self.train_boxes:\n",
        "            loc_p = loc_data[pos_idx].view(-1, 4)\n",
        "            loc_t = loc_t[pos_idx].view(-1, 4)\n",
        "            losses['B'] = F.smooth_l1_loss(loc_p, loc_t, reduction='sum') * self.bbox_alpha\n",
        "\n",
        "        if self.train_masks:\n",
        "            if self.mask_type == mask_type.direct:\n",
        "                if self.use_gt_bboxes:\n",
        "                    pos_masks = []\n",
        "                    for idx in range(batch_size):\n",
        "                        pos_masks.append(masks[idx][idx_t[idx, pos[idx]]])\n",
        "                    masks_t = torch.cat(pos_masks, 0)\n",
        "                    masks_p = mask_data[pos, :].view(-1, self.mask_dim)\n",
        "                    losses['M'] = F.binary_cross_entropy(torch.clamp(masks_p, 0, 1), masks_t, reduction='sum') * self.mask_alpha\n",
        "                else:\n",
        "                    losses['M'] = self.direct_mask_loss(pos_idx, idx_t, loc_data, mask_data, priors, masks)\n",
        "            elif self.mask_type == mask_type.lincomb:\n",
        "                ret = self.lincomb_mask_loss(pos, idx_t, loc_data, mask_data, priors, proto_data, masks, gt_box_t, score_data, inst_data, labels)\n",
        "                if self.use_maskiou:\n",
        "                    loss, maskiou_targets = ret\n",
        "                else:\n",
        "                    loss = ret\n",
        "                losses.update(loss)\n",
        "\n",
        "                if self.mask_proto_loss is not None:\n",
        "                    if self.mask_proto_loss == 'l1':\n",
        "                        losses['P'] = torch.mean(torch.abs(proto_data)) / self.l1_expected_area * self.l1_alpha\n",
        "                    elif self.mask_proto_loss == 'disj':\n",
        "                        losses['P'] = -torch.mean(torch.max(F.log_softmax(proto_data, dim=-1), dim=-1)[0])\n",
        "\n",
        "        # Confidence loss\n",
        "        if self.use_focal_loss:\n",
        "            if self.use_sigmoid_focal_loss:\n",
        "                losses['C'] = self.focal_conf_sigmoid_loss(conf_data, conf_t)\n",
        "            elif self.use_objectness_score:\n",
        "                losses['C'] = self.focal_conf_objectness_loss(conf_data, conf_t)\n",
        "            else:\n",
        "                losses['C'] = self.focal_conf_loss(conf_data, conf_t)\n",
        "        else:\n",
        "            if self.use_objectness_score:\n",
        "                losses['C'] = self.conf_objectness_loss(conf_data, conf_t, batch_size, loc_p, loc_t, priors)\n",
        "            else:\n",
        "                losses['C'] = self.ohem_conf_loss(conf_data, conf_t, pos, batch_size)\n",
        "\n",
        "        # Mask IoU Loss\n",
        "        if self.use_maskiou and maskiou_targets is not None:\n",
        "            losses['I'] = self.mask_iou_loss(net, maskiou_targets)\n",
        "\n",
        "        # These losses also don't depend on anchors\n",
        "        if self.use_class_existence_loss:\n",
        "            losses['E'] = self.class_existence_loss(predictions['classes'], class_existence_t)\n",
        "        if self.use_semantic_segmentation_loss:\n",
        "            losses['S'] = self.semantic_segmentation_loss(predictions['segm'], masks, labels)\n",
        "\n",
        "        # Divide all losses by the number of positives.\n",
        "        # Don't do it for loss[P] because that doesn't depend on the anchors.\n",
        "        total_num_pos = num_pos.data.sum().float()\n",
        "        for k in losses:\n",
        "            if k not in ('P', 'E', 'S'):\n",
        "                losses[k] /= total_num_pos\n",
        "            else:\n",
        "                losses[k] /= batch_size\n",
        "\n",
        "        # Loss Key:\n",
        "        #  - B: Box Localization Loss\n",
        "        #  - C: Class Confidence Loss\n",
        "        #  - M: Mask Loss\n",
        "        #  - P: Prototype Loss\n",
        "        #  - D: Coefficient Diversity Loss\n",
        "        #  - E: Class Existence Loss\n",
        "        #  - S: Semantic Segmentation Loss\n",
        "        return losses\n",
        "\n",
        "    def class_existence_loss(self, class_data, class_existence_t):\n",
        "        return self.class_existence_alpha * F.binary_cross_entropy_with_logits(class_data, class_existence_t, reduction='sum')\n",
        "\n",
        "    def semantic_segmentation_loss(self, segment_data, mask_t, class_t, interpolation_mode='bilinear'):\n",
        "        # Note num_classes here is without the background class so self.num_classes-1\n",
        "        batch_size, num_classes, mask_h, mask_w = segment_data.size()\n",
        "        loss_s = 0\n",
        "\n",
        "        for idx in range(batch_size):\n",
        "            cur_segment = segment_data[idx]\n",
        "            cur_class_t = class_t[idx]\n",
        "\n",
        "            with torch.no_grad():\n",
        "                downsampled_masks = F.interpolate(mask_t[idx].unsqueeze(0), (mask_h, mask_w),\n",
        "                                                  mode=interpolation_mode, align_corners=False).squeeze(0)\n",
        "                downsampled_masks = downsampled_masks.gt(0.5).float()\n",
        "\n",
        "                # Construct Semantic Segmentation\n",
        "                segment_t = torch.zeros_like(cur_segment, requires_grad=False)\n",
        "                for obj_idx in range(downsampled_masks.size(0)):\n",
        "                    segment_t[cur_class_t[obj_idx]] = torch.max(segment_t[cur_class_t[obj_idx]], downsampled_masks[obj_idx])\n",
        "\n",
        "            loss_s += F.binary_cross_entropy_with_logits(cur_segment, segment_t, reduction='sum')\n",
        "\n",
        "        return loss_s / mask_h / mask_w * self.semantic_segmentation_alpha\n",
        "\n",
        "\n",
        "    def ohem_conf_loss(self, conf_data, conf_t, pos, num):\n",
        "        # Compute max conf across batch for hard negative mining\n",
        "        batch_conf = conf_data.view(-1, self.num_classes)\n",
        "        if self.ohem_use_most_confident:\n",
        "            # i.e. max(softmax) along classes > 0\n",
        "            batch_conf = F.softmax(batch_conf, dim=1)\n",
        "            loss_c, _ = batch_conf[:, 1:].max(dim=1)\n",
        "        else:\n",
        "            # i.e. -softmax(class 0 confidence)\n",
        "            loss_c = log_sum_exp(batch_conf) - batch_conf[:, 0]\n",
        "\n",
        "        # Hard Negative Mining\n",
        "        loss_c = loss_c.view(num, -1)\n",
        "        loss_c[pos]        = 0 # filter out pos boxes\n",
        "        loss_c[conf_t < 0] = 0 # filter out neutrals (conf_t = -1)\n",
        "        _, loss_idx = loss_c.sort(1, descending=True)\n",
        "        _, idx_rank = loss_idx.sort(1)\n",
        "        num_pos = pos.long().sum(1, keepdim=True)\n",
        "        num_neg = torch.clamp(self.negpos_ratio*num_pos, max=pos.size(1)-1)\n",
        "        neg = idx_rank < num_neg.expand_as(idx_rank)\n",
        "\n",
        "        # Just in case there aren't enough negatives, don't start using positives as negatives\n",
        "        neg[pos]        = 0\n",
        "        neg[conf_t < 0] = 0 # Filter out neutrals\n",
        "\n",
        "        # Confidence Loss Including Positive and Negative Examples\n",
        "        pos_idx = pos.unsqueeze(2).expand_as(conf_data)\n",
        "        neg_idx = neg.unsqueeze(2).expand_as(conf_data)\n",
        "        conf_p = conf_data[(pos_idx+neg_idx).gt(0)].view(-1, self.num_classes)\n",
        "        targets_weighted = conf_t[(pos+neg).gt(0)]\n",
        "        loss_c = F.cross_entropy(conf_p, targets_weighted, reduction='none')\n",
        "\n",
        "        if self.use_class_balanced_conf:\n",
        "            # Lazy initialization\n",
        "            if self.class_instances is None:\n",
        "                self.class_instances = torch.zeros(self.num_classes, device=targets_weighted.device)\n",
        "\n",
        "            classes, counts = targets_weighted.unique(return_counts=True)\n",
        "\n",
        "            for _cls, _cnt in zip(classes.cpu().numpy(), counts.cpu().numpy()):\n",
        "                self.class_instances[_cls] += _cnt\n",
        "\n",
        "            self.total_instances += targets_weighted.size(0)\n",
        "\n",
        "            weighting = 1 - (self.class_instances[targets_weighted] / self.total_instances)\n",
        "            weighting = torch.clamp(weighting, min=1/self.num_classes)\n",
        "\n",
        "            # If you do the math, the average weight of self.class_instances is this\n",
        "            avg_weight = (self.num_classes - 1) / self.num_classes\n",
        "\n",
        "            loss_c = (loss_c * weighting).sum() / avg_weight\n",
        "        else:\n",
        "            loss_c = loss_c.sum()\n",
        "\n",
        "        return self.conf_alpha * loss_c\n",
        "\n",
        "    def focal_conf_loss(self, conf_data, conf_t):\n",
        "        \"\"\"\n",
        "        Focal loss as described in https://arxiv.org/pdf/1708.02002.pdf\n",
        "        Adapted from https://github.com/clcarwin/focal_loss_pytorch/blob/master/focalloss.py\n",
        "        Note that this uses softmax and not the original sigmoid from the paper.\n",
        "        \"\"\"\n",
        "        conf_t = conf_t.view(-1) # [batch_size*num_priors]\n",
        "        conf_data = conf_data.view(-1, conf_data.size(-1)) # [batch_size*num_priors, num_classes]\n",
        "\n",
        "        # Ignore neutral samples (class < 0)\n",
        "        keep = (conf_t >= 0).float()\n",
        "        conf_t[conf_t < 0] = 0 # so that gather doesn't drum up a fuss\n",
        "\n",
        "        logpt = F.log_softmax(conf_data, dim=-1)\n",
        "        logpt = logpt.gather(1, conf_t.unsqueeze(-1))\n",
        "        logpt = logpt.view(-1)\n",
        "        pt    = logpt.exp()\n",
        "\n",
        "        # I adapted the alpha_t calculation here from\n",
        "        # https://github.com/pytorch/pytorch/blob/master/modules/detectron/softmax_focal_loss_op.cu\n",
        "        # You'd think you want all the alphas to sum to one, but in the original implementation they\n",
        "        # just give background an alpha of 1-alpha and each forground an alpha of alpha.\n",
        "        background = (conf_t == 0).float()\n",
        "        at = (1 - self.focal_loss_alpha) * background + self.focal_loss_alpha * (1 - background)\n",
        "\n",
        "        loss = -at * (1 - pt) ** self.focal_loss_gamma * logpt\n",
        "\n",
        "        # See comment above for keep\n",
        "        return self.conf_alpha * (loss * keep).sum()\n",
        "\n",
        "    def focal_conf_sigmoid_loss(self, conf_data, conf_t):\n",
        "        \"\"\"\n",
        "        Focal loss but using sigmoid like the original paper.\n",
        "        Note: To make things mesh easier, the network still predicts 81 class confidences in this mode.\n",
        "              Because retinanet originally only predicts 80, we simply just don't use conf_data[..., 0]\n",
        "        \"\"\"\n",
        "        num_classes = conf_data.size(-1)\n",
        "\n",
        "        conf_t = conf_t.view(-1) # [batch_size*num_priors]\n",
        "        conf_data = conf_data.view(-1, num_classes) # [batch_size*num_priors, num_classes]\n",
        "\n",
        "        # Ignore neutral samples (class < 0)\n",
        "        keep = (conf_t >= 0).float()\n",
        "        conf_t[conf_t < 0] = 0 # can't mask with -1, so filter that out\n",
        "\n",
        "        # Compute a one-hot embedding of conf_t\n",
        "        # From https://github.com/kuangliu/pytorch-retinanet/blob/master/utils.py\n",
        "        conf_one_t = torch.eye(num_classes, device=conf_t.get_device())[conf_t]\n",
        "        conf_pm_t  = conf_one_t * 2 - 1 # -1 if background, +1 if forground for specific class\n",
        "\n",
        "        logpt = F.logsigmoid(conf_data * conf_pm_t) # note: 1 - sigmoid(x) = sigmoid(-x)\n",
        "        pt    = logpt.exp()\n",
        "\n",
        "        at = self.focal_loss_alpha * conf_one_t + (1 - self.focal_loss_alpha) * (1 - conf_one_t)\n",
        "        at[..., 0] = 0 # Set alpha for the background class to 0 because sigmoid focal loss doesn't use it\n",
        "\n",
        "        loss = -at * (1 - pt) ** self.focal_loss_gamma * logpt\n",
        "        loss = keep * loss.sum(dim=-1)\n",
        "\n",
        "        return self.conf_alpha * loss.sum()\n",
        "\n",
        "    def focal_conf_objectness_loss(self, conf_data, conf_t):\n",
        "        \"\"\"\n",
        "        Instead of using softmax, use class[0] to be the objectness score and do sigmoid focal loss on that.\n",
        "        Then for the rest of the classes, softmax them and apply CE for only the positive examples.\n",
        "\n",
        "        If class[0] = 1 implies forground and class[0] = 0 implies background then you achieve something\n",
        "        similar during test-time to softmax by setting class[1:] = softmax(class[1:]) * class[0] and invert class[0].\n",
        "        \"\"\"\n",
        "\n",
        "        conf_t = conf_t.view(-1) # [batch_size*num_priors]\n",
        "        conf_data = conf_data.view(-1, conf_data.size(-1)) # [batch_size*num_priors, num_classes]\n",
        "\n",
        "        # Ignore neutral samples (class < 0)\n",
        "        keep = (conf_t >= 0).float()\n",
        "        conf_t[conf_t < 0] = 0 # so that gather doesn't drum up a fuss\n",
        "\n",
        "        background = (conf_t == 0).float()\n",
        "        at = (1 - self.focal_loss_alpha) * background + self.focal_loss_alpha * (1 - background)\n",
        "\n",
        "        logpt = F.logsigmoid(conf_data[:, 0]) * (1 - background) + F.logsigmoid(-conf_data[:, 0]) * background\n",
        "        pt    = logpt.exp()\n",
        "\n",
        "        obj_loss = -at * (1 - pt) ** self.focal_loss_gamma * logpt\n",
        "\n",
        "        # All that was the objectiveness loss--now time for the class confidence loss\n",
        "        pos_mask = conf_t > 0\n",
        "        conf_data_pos = (conf_data[:, 1:])[pos_mask] # Now this has just 80 classes\n",
        "        conf_t_pos    = conf_t[pos_mask] - 1         # So subtract 1 here\n",
        "\n",
        "        class_loss = F.cross_entropy(conf_data_pos, conf_t_pos, reduction='sum')\n",
        "\n",
        "        return self.conf_alpha * (class_loss + (obj_loss * keep).sum())\n",
        "\n",
        "    def conf_objectness_loss(self, conf_data, conf_t, batch_size, loc_p, loc_t, priors):\n",
        "        \"\"\"\n",
        "        Instead of using softmax, use class[0] to be p(obj) * p(IoU) as in YOLO.\n",
        "        Then for the rest of the classes, softmax them and apply CE for only the positive examples.\n",
        "        \"\"\"\n",
        "\n",
        "        conf_t = conf_t.view(-1) # [batch_size*num_priors]\n",
        "        conf_data = conf_data.view(-1, conf_data.size(-1)) # [batch_size*num_priors, num_classes]\n",
        "\n",
        "        pos_mask = (conf_t > 0)\n",
        "        neg_mask = (conf_t == 0)\n",
        "\n",
        "        obj_data = conf_data[:, 0]\n",
        "        obj_data_pos = obj_data[pos_mask]\n",
        "        obj_data_neg = obj_data[neg_mask]\n",
        "\n",
        "        # Don't be confused, this is just binary cross entropy similified\n",
        "        obj_neg_loss = - F.logsigmoid(-obj_data_neg).sum()\n",
        "\n",
        "        with torch.no_grad():\n",
        "            pos_priors = priors.unsqueeze(0).expand(batch_size, -1, -1).reshape(-1, 4)[pos_mask, :]\n",
        "\n",
        "            boxes_pred = decode(loc_p, pos_priors, self.use_yolo_regressors)\n",
        "            boxes_targ = decode(loc_t, pos_priors, self.use_yolo_regressors)\n",
        "\n",
        "            iou_targets = elemwise_box_iou(boxes_pred, boxes_targ)\n",
        "\n",
        "        obj_pos_loss = - iou_targets * F.logsigmoid(obj_data_pos) - (1 - iou_targets) * F.logsigmoid(-obj_data_pos)\n",
        "        obj_pos_loss = obj_pos_loss.sum()\n",
        "\n",
        "        # All that was the objectiveness loss--now time for the class confidence loss\n",
        "        conf_data_pos = (conf_data[:, 1:])[pos_mask] # Now this has just 80 classes\n",
        "        conf_t_pos    = conf_t[pos_mask] - 1         # So subtract 1 here\n",
        "\n",
        "        class_loss = F.cross_entropy(conf_data_pos, conf_t_pos, reduction='sum')\n",
        "\n",
        "        return self.conf_alpha * (class_loss + obj_pos_loss + obj_neg_loss)\n",
        "\n",
        "\n",
        "    def direct_mask_loss(self, pos_idx, idx_t, loc_data, mask_data, priors, masks):\n",
        "        \"\"\" Crops the gt masks using the predicted bboxes, scales them down, and outputs the BCE loss. \"\"\"\n",
        "        loss_m = 0\n",
        "        for idx in range(mask_data.size(0)):\n",
        "            with torch.no_grad():\n",
        "                cur_pos_idx = pos_idx[idx, :, :]\n",
        "                cur_pos_idx_squeezed = cur_pos_idx[:, 1]\n",
        "\n",
        "                # Shape: [num_priors, 4], decoded predicted bboxes\n",
        "                pos_bboxes = decode(loc_data[idx, :, :], priors.data, self.use_yolo_regressors)\n",
        "                pos_bboxes = pos_bboxes[cur_pos_idx].view(-1, 4).clamp(0, 1)\n",
        "                pos_lookup = idx_t[idx, cur_pos_idx_squeezed]\n",
        "\n",
        "                cur_masks = masks[idx]\n",
        "                pos_masks = cur_masks[pos_lookup, :, :]\n",
        "\n",
        "                # Convert bboxes to absolute coordinates\n",
        "                num_pos, img_height, img_width = pos_masks.size()\n",
        "\n",
        "                # Take care of all the bad behavior that can be caused by out of bounds coordinates\n",
        "                x1, x2 = sanitize_coordinates(pos_bboxes[:, 0], pos_bboxes[:, 2], img_width)\n",
        "                y1, y2 = sanitize_coordinates(pos_bboxes[:, 1], pos_bboxes[:, 3], img_height)\n",
        "\n",
        "                # Crop each gt mask with the predicted bbox and rescale to the predicted mask size\n",
        "                # Note that each bounding box crop is a different size so I don't think we can vectorize this\n",
        "                scaled_masks = []\n",
        "                for jdx in range(num_pos):\n",
        "                    tmp_mask = pos_masks[jdx, y1[jdx]:y2[jdx], x1[jdx]:x2[jdx]]\n",
        "\n",
        "                    # Restore any dimensions we've left out because our bbox was 1px wide\n",
        "                    while tmp_mask.dim() < 2:\n",
        "                        tmp_mask = tmp_mask.unsqueeze(0)\n",
        "\n",
        "                    new_mask = F.adaptive_avg_pool2d(tmp_mask.unsqueeze(0), self.mask_size)\n",
        "                    scaled_masks.append(new_mask.view(1, -1))\n",
        "\n",
        "                mask_t = torch.cat(scaled_masks, 0).gt(0.5).float() # Threshold downsampled mask\n",
        "\n",
        "            pos_mask_data = mask_data[idx, cur_pos_idx_squeezed, :]\n",
        "            loss_m += F.binary_cross_entropy(torch.clamp(pos_mask_data, 0, 1), mask_t, reduction='sum') * self.mask_alpha\n",
        "\n",
        "        return loss_m\n",
        "\n",
        "\n",
        "    def coeff_diversity_loss(self, coeffs, instance_t):\n",
        "        \"\"\"\n",
        "        coeffs     should be size [num_pos, num_coeffs]\n",
        "        instance_t should be size [num_pos] and be values from 0 to num_instances-1\n",
        "        \"\"\"\n",
        "        num_pos = coeffs.size(0)\n",
        "        instance_t = instance_t.view(-1) # juuuust to make sure\n",
        "\n",
        "        coeffs_norm = F.normalize(coeffs, dim=1)\n",
        "        cos_sim = coeffs_norm @ coeffs_norm.t()\n",
        "\n",
        "        inst_eq = (instance_t[:, None].expand_as(cos_sim) == instance_t[None, :].expand_as(cos_sim)).float()\n",
        "\n",
        "        # Rescale to be between 0 and 1\n",
        "        cos_sim = (cos_sim + 1) / 2\n",
        "\n",
        "        # If they're the same instance, use cosine distance, else use cosine similarity\n",
        "        loss = (1 - cos_sim) * inst_eq + cos_sim * (1 - inst_eq)\n",
        "\n",
        "        # Only divide by num_pos once because we're summing over a num_pos x num_pos tensor\n",
        "        # and all the losses will be divided by num_pos at the end, so just one extra time.\n",
        "        return self.mask_proto_coeff_diversity_alpha * loss.sum() / num_pos\n",
        "\n",
        "\n",
        "    def lincomb_mask_loss(self, pos, idx_t, loc_data, mask_data, priors, proto_data, masks, gt_box_t, score_data, inst_data, labels, interpolation_mode='bilinear'):\n",
        "        mask_h = proto_data.size(1)\n",
        "        mask_w = proto_data.size(2)\n",
        "\n",
        "        process_gt_bboxes = self.mask_proto_normalize_emulate_roi_pooling or self.mask_proto_crop\n",
        "\n",
        "        if self.mask_proto_remove_empty_masks:\n",
        "            # Make sure to store a copy of this because we edit it to get rid of all-zero masks\n",
        "            pos = pos.clone()\n",
        "\n",
        "        loss_m = 0\n",
        "        loss_d = 0 # Coefficient diversity loss\n",
        "\n",
        "        maskiou_t_list = []\n",
        "        maskiou_net_input_list = []\n",
        "        label_t_list = []\n",
        "\n",
        "        for idx in range(mask_data.size(0)):\n",
        "            with torch.no_grad():\n",
        "                downsampled_masks = F.interpolate(masks[idx].unsqueeze(0), (mask_h, mask_w),\n",
        "                                                  mode=interpolation_mode, align_corners=False).squeeze(0)\n",
        "                downsampled_masks = downsampled_masks.permute(1, 2, 0).contiguous()\n",
        "\n",
        "                if self.mask_proto_binarize_downsampled_gt:\n",
        "                    downsampled_masks = downsampled_masks.gt(0.5).float()\n",
        "\n",
        "                if self.mask_proto_remove_empty_masks:\n",
        "                    # Get rid of gt masks that are so small they get downsampled away\n",
        "                    very_small_masks = (downsampled_masks.sum(dim=(0,1)) <= 0.0001)\n",
        "                    for i in range(very_small_masks.size(0)):\n",
        "                        if very_small_masks[i]:\n",
        "                            pos[idx, idx_t[idx] == i] = 0\n",
        "\n",
        "                if self.mask_proto_reweight_mask_loss:\n",
        "                    # Ensure that the gt is binary\n",
        "                    if not self.mask_proto_binarize_downsampled_gt:\n",
        "                        bin_gt = downsampled_masks.gt(0.5).float()\n",
        "                    else:\n",
        "                        bin_gt = downsampled_masks\n",
        "\n",
        "                    gt_foreground_norm = bin_gt     / (torch.sum(bin_gt,   dim=(0,1), keepdim=True) + 0.0001)\n",
        "                    gt_background_norm = (1-bin_gt) / (torch.sum(1-bin_gt, dim=(0,1), keepdim=True) + 0.0001)\n",
        "\n",
        "                    mask_reweighting   = gt_foreground_norm * self.mask_proto_reweight_coeff + gt_background_norm\n",
        "                    mask_reweighting  *= mask_h * mask_w\n",
        "\n",
        "            cur_pos = pos[idx]\n",
        "            pos_idx_t = idx_t[idx, cur_pos]\n",
        "\n",
        "            if process_gt_bboxes:\n",
        "                # Note: this is in point-form\n",
        "                if self.mask_proto_crop_with_pred_box:\n",
        "                    pos_gt_box_t = decode(loc_data[idx, :, :], priors.data, self.use_yolo_regressors)[cur_pos]\n",
        "                else:\n",
        "                    pos_gt_box_t = gt_box_t[idx, cur_pos]\n",
        "\n",
        "            if pos_idx_t.size(0) == 0:\n",
        "                continue\n",
        "\n",
        "            proto_masks = proto_data[idx]\n",
        "            proto_coef  = mask_data[idx, cur_pos, :]\n",
        "            if self.use_mask_scoring:\n",
        "                mask_scores = score_data[idx, cur_pos, :]\n",
        "\n",
        "            if self.mask_proto_coeff_diversity_loss:\n",
        "                if inst_data is not None:\n",
        "                    div_coeffs = inst_data[idx, cur_pos, :]\n",
        "                else:\n",
        "                    div_coeffs = proto_coef\n",
        "\n",
        "                loss_d += self.coeff_diversity_loss(div_coeffs, pos_idx_t)\n",
        "\n",
        "            # If we have over the allowed number of masks, select a random sample\n",
        "            old_num_pos = proto_coef.size(0)\n",
        "            if old_num_pos > self.masks_to_train:\n",
        "                perm = torch.randperm(proto_coef.size(0))\n",
        "                select = perm[:self.masks_to_train]\n",
        "\n",
        "                proto_coef = proto_coef[select, :]\n",
        "                pos_idx_t  = pos_idx_t[select]\n",
        "\n",
        "                if process_gt_bboxes:\n",
        "                    pos_gt_box_t = pos_gt_box_t[select, :]\n",
        "                if self.use_mask_scoring:\n",
        "                    mask_scores = mask_scores[select, :]\n",
        "\n",
        "            num_pos = proto_coef.size(0)\n",
        "            mask_t = downsampled_masks[:, :, pos_idx_t]\n",
        "            label_t = labels[idx][pos_idx_t]\n",
        "\n",
        "            # Size: [mask_h, mask_w, num_pos]\n",
        "            pred_masks = proto_masks @ proto_coef.t()\n",
        "            pred_masks = self.mask_proto_mask_activation(pred_masks)\n",
        "\n",
        "            if self.mask_proto_double_loss:\n",
        "                if self.mask_proto_mask_activation == activation_func[\"sigmoid\"]:\n",
        "                    pre_loss = F.binary_cross_entropy(torch.clamp(pred_masks, 0, 1), mask_t, reduction='sum')\n",
        "                else:\n",
        "                    pre_loss = F.smooth_l1_loss(pred_masks, mask_t, reduction='sum')\n",
        "\n",
        "                loss_m += self.mask_proto_double_loss_alpha * pre_loss\n",
        "\n",
        "            if self.mask_proto_crop:\n",
        "                pred_masks = crop(pred_masks, pos_gt_box_t)\n",
        "\n",
        "            if self.mask_proto_mask_activation == activation_func[\"sigmoid\"]:\n",
        "                pre_loss = F.binary_cross_entropy(torch.clamp(pred_masks, 0, 1), mask_t, reduction='none')\n",
        "            else:\n",
        "                pre_loss = F.smooth_l1_loss(pred_masks, mask_t, reduction='none')\n",
        "\n",
        "            if self.mask_proto_normalize_mask_loss_by_sqrt_area:\n",
        "                gt_area  = torch.sum(mask_t, dim=(0, 1), keepdim=True)\n",
        "                pre_loss = pre_loss / (torch.sqrt(gt_area) + 0.0001)\n",
        "\n",
        "            if self.mask_proto_reweight_mask_loss:\n",
        "                pre_loss = pre_loss * mask_reweighting[:, :, pos_idx_t]\n",
        "\n",
        "            if self.mask_proto_normalize_emulate_roi_pooling:\n",
        "                weight = mask_h * mask_w if self.mask_proto_crop else 1\n",
        "                pos_gt_csize = center_size(pos_gt_box_t)\n",
        "                gt_box_width  = pos_gt_csize[:, 2] * mask_w\n",
        "                gt_box_height = pos_gt_csize[:, 3] * mask_h\n",
        "                pre_loss = pre_loss.sum(dim=(0, 1)) / gt_box_width / gt_box_height * weight\n",
        "\n",
        "            # If the number of masks were limited scale the loss accordingly\n",
        "            if old_num_pos > num_pos:\n",
        "                pre_loss *= old_num_pos / num_pos\n",
        "\n",
        "            loss_m += torch.sum(pre_loss)\n",
        "\n",
        "            if self.use_maskiou:\n",
        "                if self.discard_mask_area > 0:\n",
        "                    gt_mask_area = torch.sum(mask_t, dim=(0, 1))\n",
        "                    select = gt_mask_area > self.discard_mask_area\n",
        "\n",
        "                    if torch.sum(select) < 1:\n",
        "                        continue\n",
        "\n",
        "                    pos_gt_box_t = pos_gt_box_t[select, :]\n",
        "                    pred_masks = pred_masks[:, :, select]\n",
        "                    mask_t = mask_t[:, :, select]\n",
        "                    label_t = label_t[select]\n",
        "\n",
        "                maskiou_net_input = pred_masks.permute(2, 0, 1).contiguous().unsqueeze(1)\n",
        "                pred_masks = pred_masks.gt(0.5).float()\n",
        "                maskiou_t = self._mask_iou(pred_masks, mask_t)\n",
        "\n",
        "                maskiou_net_input_list.append(maskiou_net_input)\n",
        "                maskiou_t_list.append(maskiou_t)\n",
        "                label_t_list.append(label_t)\n",
        "\n",
        "        losses = {'M': loss_m * self.mask_alpha / mask_h / mask_w}\n",
        "\n",
        "        if self.mask_proto_coeff_diversity_loss:\n",
        "            losses['D'] = loss_d\n",
        "\n",
        "        if self.use_maskiou:\n",
        "            # discard_mask_area discarded every mask in the batch, so nothing to do here\n",
        "            if len(maskiou_t_list) == 0:\n",
        "                return losses, None\n",
        "\n",
        "            maskiou_t = torch.cat(maskiou_t_list)\n",
        "            label_t = torch.cat(label_t_list)\n",
        "            maskiou_net_input = torch.cat(maskiou_net_input_list)\n",
        "\n",
        "            num_samples = maskiou_t.size(0)\n",
        "            if self.maskious_to_train > 0 and num_samples > self.maskious_to_train:\n",
        "                perm = torch.randperm(num_samples)\n",
        "                select = perm[:self.masks_to_train]\n",
        "                maskiou_t = maskiou_t[select]\n",
        "                label_t = label_t[select]\n",
        "                maskiou_net_input = maskiou_net_input[select]\n",
        "\n",
        "            return losses, [maskiou_net_input, maskiou_t, label_t]\n",
        "\n",
        "        return losses\n",
        "\n",
        "    def _mask_iou(self, mask1, mask2):\n",
        "        intersection = torch.sum(mask1*mask2, dim=(0, 1))\n",
        "        area1 = torch.sum(mask1, dim=(0, 1))\n",
        "        area2 = torch.sum(mask2, dim=(0, 1))\n",
        "        union = (area1 + area2) - intersection\n",
        "        ret = intersection / union\n",
        "        return ret\n",
        "\n",
        "    def mask_iou_loss(self, net, maskiou_targets):\n",
        "        maskiou_net_input, maskiou_t, label_t = maskiou_targets\n",
        "\n",
        "        maskiou_p = net.maskiou_net(maskiou_net_input)\n",
        "\n",
        "        label_t = label_t[:, None]\n",
        "        maskiou_p = torch.gather(maskiou_p, dim=1, index=label_t).view(-1)\n",
        "\n",
        "        loss_i = F.smooth_l1_loss(maskiou_p, maskiou_t, reduction='sum')\n",
        "\n",
        "        return loss_i * self.maskiou_alpha"
      ],
      "metadata": {
        "id": "NHQ9TvPLHpe7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "ahMnvCxLFOUT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oErx-J9A7-Gm",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "341801ca-0de8-4c0d-cbf8-165095765cc6"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.9/7.9 MB\u001b[0m \u001b[31m14.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m42.6/42.6 kB\u001b[0m \u001b[31m4.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m139.3/139.3 kB\u001b[0m \u001b[31m20.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m54.4/54.4 kB\u001b[0m \u001b[31m8.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m59.9/59.9 kB\u001b[0m \u001b[31m7.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m79.9/79.9 MB\u001b[0m \u001b[31m21.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m108.9/108.9 kB\u001b[0m \u001b[31m14.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m74.3/74.3 kB\u001b[0m \u001b[31m11.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m676.9/676.9 kB\u001b[0m \u001b[31m53.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m71.9/71.9 kB\u001b[0m \u001b[31m8.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m192.5/192.5 kB\u001b[0m \u001b[31m25.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m89.6/89.6 kB\u001b[0m \u001b[31m12.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m942.9/942.9 kB\u001b[0m \u001b[31m66.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m202.9/202.9 kB\u001b[0m \u001b[31m25.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.3/58.3 kB\u001b[0m \u001b[31m7.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m57.5/57.5 kB\u001b[0m \u001b[31m7.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m307.7/307.7 kB\u001b[0m \u001b[31m36.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m75.6/75.6 kB\u001b[0m \u001b[31m9.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m116.3/116.3 kB\u001b[0m \u001b[31m15.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m67.6/67.6 kB\u001b[0m \u001b[31m10.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m12.2/12.2 MB\u001b[0m \u001b[31m97.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m82.2/82.2 kB\u001b[0m \u001b[31m11.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m77.9/77.9 kB\u001b[0m \u001b[31m11.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.1/2.1 MB\u001b[0m \u001b[31m88.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m411.2/411.2 kB\u001b[0m \u001b[31m41.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m138.9/138.9 kB\u001b[0m \u001b[31m18.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m49.7/49.7 kB\u001b[0m \u001b[31m6.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m93.1/93.1 kB\u001b[0m \u001b[31m13.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.0/3.0 MB\u001b[0m \u001b[31m84.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Building wheel for fiftyone-db (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "Migrating database to v0.23.8\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:fiftyone.migrations.runner:Migrating database to v0.23.8\n"
          ]
        }
      ],
      "source": [
        "!pip install -q fiftyone\n",
        "\n",
        "import fiftyone.zoo as foz\n",
        "from typing import Tuple\n",
        "from PIL import Image\n",
        "import torch\n",
        "\n",
        "class COCODataset(torch.utils.data.Dataset):\n",
        "\n",
        "    def __init__(\n",
        "            self,\n",
        "            split: str = \"train\",\n",
        "            max_samples: int = 100,\n",
        "            dataset_name: str = \"coco-2017\",\n",
        "            transform: torch.nn.Module = None):\n",
        "        \"\"\"\n",
        "        Initialize COCO dataset from fiftyone zoo\n",
        "\n",
        "        Args:\n",
        "            split: `train` or `test` or `validation`\n",
        "            max_samples: maximum number of samples to load\n",
        "            dataset_name: name of the dataset\n",
        "            transform: transforms to apply to the dataset\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.dataset = foz.load_zoo_dataset(\n",
        "            \"coco-2017\",\n",
        "            split=split,\n",
        "            max_samples=max_samples,\n",
        "            dataset_name=dataset_name,\n",
        "            label_types=[\"detections\", \"segmentations\"],\n",
        "            persistent=True,\n",
        "        )\n",
        "        self.label_map = {label: index for index, label in\n",
        "                          enumerate(self.dataset.default_classes)}\n",
        "        self.transform = transform\n",
        "\n",
        "\n",
        "    def __len__(self) -> int:\n",
        "        \"\"\"\n",
        "        Get the size of the dataset\n",
        "\n",
        "        Returns:\n",
        "            int: number of samples in the dataset\n",
        "        \"\"\"\n",
        "        return len(self.dataset)\n",
        "\n",
        "    def __getitem__(self, index: int) -> Tuple[torch.Tensor, tuple]:\n",
        "        \"\"\"\n",
        "        Get sample from the dataset identified by `index`\n",
        "\n",
        "        Args:\n",
        "            index: index of the sample to get\n",
        "\n",
        "        Returns:\n",
        "            tuple: image tensor, (bounding boxes, masks, num_crowds)\n",
        "        \"\"\"\n",
        "        sample =  [sample for sample in self.dataset.iter_samples()][index]\n",
        "\n",
        "        image = Image.open(sample.filepath)\n",
        "\n",
        "        if self.transform is not None:\n",
        "            image = self.transform(image)\n",
        "\n",
        "        detections = sample.segmentations.detections\n",
        "\n",
        "        height, width = image.shape[1], image.shape[2]\n",
        "\n",
        "        num_crowds = torch.zeros((1,), dtype=torch.int64)\n",
        "        # labels = []\n",
        "        bboxes = []\n",
        "        masks = torch.zeros((len(detections), height, width))\n",
        "\n",
        "        for i, detection in enumerate(detections):\n",
        "            num_crowds += detection.iscrowd\n",
        "\n",
        "            # labels.append(self.label_map[detection.label])\n",
        "\n",
        "            # top-left coordinates\n",
        "            x1, y1, w, h = detection.bounding_box\n",
        "            x1 = int(x1 * width)\n",
        "            y1 = int(y1 * height)\n",
        "            w = int(w * width)\n",
        "            h = int(h * height)\n",
        "            # bottom-right coordinates\n",
        "            x2 = x1 + w\n",
        "            y2 = y1 + h\n",
        "\n",
        "            # bbox format: x_min, y_min, x_max, y_max, label_id\n",
        "            bboxes.append([x1, y1, x2, y2, self.label_map[detection.label]])\n",
        "\n",
        "            if detection.mask is not None:\n",
        "                mask_size = (h, w)\n",
        "                mask = torch.from_numpy(\n",
        "                    detection.mask).unsqueeze(0).unsqueeze(0)\n",
        "                resized_mask = torch.nn.functional.interpolate(\n",
        "                    mask.float(),\n",
        "                    size=mask_size,\n",
        "                    mode=\"nearest\"\n",
        "                ).squeeze().byte()\n",
        "\n",
        "                masks[i, y1:y2, x1:x2] = resized_mask\n",
        "\n",
        "        # labels = torch.as_tensor(labels, dtype=torch.int32)\n",
        "        bboxes = torch.as_tensor(bboxes, dtype=torch.float32)\n",
        "\n",
        "        target = {\n",
        "            # \"labels\": labels,\n",
        "            \"bboxes\": bboxes,\n",
        "            \"masks\": masks,\n",
        "            \"num_crowds\": num_crowds,\n",
        "        }\n",
        "\n",
        "        return image, target"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from torchvision.transforms import v2\n",
        "\n",
        "COCO_PRE_MEAN = [0.485, 0.456, 0.406]\n",
        "COCO_PRE_STD = [0.229, 0.224, 0.225]\n",
        "\n",
        "def get_transform():\n",
        "    return v2.Compose([\n",
        "        v2.Resize((512, 800)),\n",
        "        v2.ToTensor(),\n",
        "        v2.Normalize(mean=COCO_PRE_MEAN, std=COCO_PRE_STD),\n",
        "    ])"
      ],
      "metadata": {
        "id": "yyo5Ogom6wYi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset = COCODataset(transform=get_transform())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "B9ihiKaF55-S",
        "outputId": "7d13802b-ab5d-49a3-be5a-9d49f208a58c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading split 'train' to '/root/fiftyone/coco-2017/train' if necessary\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torchvision/transforms/v2/_deprecated.py:41: UserWarning: The transform `ToTensor()` is deprecated and will be removed in a future release. Instead, please use `v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])`.\n",
            "  warnings.warn(\n",
            "INFO:fiftyone.zoo.datasets:Downloading split 'train' to '/root/fiftyone/coco-2017/train' if necessary\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Found annotations at '/root/fiftyone/coco-2017/raw/instances_train2017.json'\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:fiftyone.utils.coco:Found annotations at '/root/fiftyone/coco-2017/raw/instances_train2017.json'\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Sufficient images already downloaded\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:fiftyone.utils.coco:Sufficient images already downloaded\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Existing download of split 'train' is sufficient\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:fiftyone.zoo.datasets:Existing download of split 'train' is sufficient\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loading existing dataset 'coco-2017'. To reload from disk, either delete the existing dataset or provide a custom `dataset_name` to use\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:fiftyone.zoo.datasets:Loading existing dataset 'coco-2017'. To reload from disk, either delete the existing dataset or provide a custom `dataset_name` to use\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "img, target = dataset[1]"
      ],
      "metadata": {
        "id": "7-0iDVyq6dB8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "bboxes, masks, num_crowds = target[\"bboxes\"], target[\"masks\"], target[\"num_crowds\"]"
      ],
      "metadata": {
        "id": "guKU8l43NMde"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "img.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TgGsv4CL9jem",
        "outputId": "1a88451d-bc5a-4785-a425-397a0137783d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([3, 512, 800])"
            ]
          },
          "metadata": {},
          "execution_count": 70
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "bboxes.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5CuC3sH99ytk",
        "outputId": "f38d7764-c1fd-4a54-bb1c-bdbc2b067b25"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([2, 5])"
            ]
          },
          "metadata": {},
          "execution_count": 71
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "bboxes[0]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Q3xV3rUz96o8",
        "outputId": "fa0c9ede-a4e6-4c0f-e633-edad1c6ab44d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([481.,  72., 749., 429.,  25.])"
            ]
          },
          "metadata": {},
          "execution_count": 72
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "masks.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rtZqtOsT974d",
        "outputId": "44a3f084-0871-48be-fda4-9eb04b597b13"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([2, 512, 800])"
            ]
          },
          "metadata": {},
          "execution_count": 73
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "run_this = False # @param {type: \"boolean\"}\n",
        "\n",
        "if run_this:\n",
        "    for i in range(len(dataset)):\n",
        "        try:\n",
        "            dataset[i]\n",
        "            print(f\"successfully got {i}\")\n",
        "        except:\n",
        "            print(f\"failed to get {i}\")"
      ],
      "metadata": {
        "id": "j4Hzk2zv6cpO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def collate_fn(batch):\n",
        "    images = [image for image, _ in batch]\n",
        "    targets = [target for _, target in batch]\n",
        "\n",
        "    max_detections = max(len(target[\"bboxes\"]) for target in targets)\n",
        "\n",
        "    print(max_detections)\n",
        "\n",
        "    for target in targets:\n",
        "        num_detections = len(target[\"bboxes\"])\n",
        "        num_pad = max_detections - num_detections\n",
        "        if num_pad:\n",
        "            target[\"bboxes\"] = torch.cat([target[\"bboxes\"],\n",
        "                                          torch.zeros(num_pad, 5)])\n",
        "            # target[\"labels\"] = torch.cat([\n",
        "            #     target[\"labels\"],\n",
        "            #     torch.zeros(num_pad, dtype=torch.int64)])\n",
        "            target[\"masks\"] = torch.cat([\n",
        "                target[\"masks\"],\n",
        "                torch.zeros(num_pad,\n",
        "                            target[\"masks\"].shape[1],\n",
        "                            target[\"masks\"].shape[2])])\n",
        "\n",
        "    images = torch.stack(images)\n",
        "    targets = {\n",
        "        \"bboxes\": torch.stack([target[\"bboxes\"] for target in targets]),\n",
        "        # \"labels\": torch.stack([target[\"labels\"] for target in targets]),\n",
        "        \"masks\": torch.stack([target[\"masks\"] for target in targets]),\n",
        "        \"num_crowds\": target[\"num_crowds\"],\n",
        "    }\n",
        "\n",
        "    return images, targets"
      ],
      "metadata": {
        "id": "_c42945l-jxD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_loader = torch.utils.data.DataLoader(\n",
        "    dataset, batch_size=10, shuffle=True, collate_fn=collate_fn\n",
        ")"
      ],
      "metadata": {
        "id": "Uy_um24R-crW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset = COCODataset(transform=get_transform())\n",
        "criterion = MultiBoxLoss(\n",
        "        num_classes=len(dataset.dataset.default_classes),\n",
        "        pos_threshold=0.5,\n",
        "        neg_threshold=0.5,\n",
        "        # When using ohem, the ratio between positives and negatives\n",
        "        # (3 means 3 negatives to 1 positive)\n",
        "        negpos_ratio=3,\n",
        "    )"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3a3VKPDGFXDF",
        "outputId": "0a0e67bc-3315-4575-fb38-aaa7a5bef9a6"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading split 'train' to '/root/fiftyone/coco-2017/train' if necessary\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torchvision/transforms/v2/_deprecated.py:41: UserWarning: The transform `ToTensor()` is deprecated and will be removed in a future release. Instead, please use `v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])`.\n",
            "  warnings.warn(\n",
            "INFO:fiftyone.zoo.datasets:Downloading split 'train' to '/root/fiftyone/coco-2017/train' if necessary\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading annotations to '/root/fiftyone/coco-2017/tmp-download/annotations_trainval2017.zip'\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:fiftyone.utils.coco:Downloading annotations to '/root/fiftyone/coco-2017/tmp-download/annotations_trainval2017.zip'\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            " 100% |██████|    1.9Gb/1.9Gb [17.5s elapsed, 0s remaining, 122.6Mb/s]      \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:eta.core.utils: 100% |██████|    1.9Gb/1.9Gb [17.5s elapsed, 0s remaining, 122.6Mb/s]      \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting annotations to '/root/fiftyone/coco-2017/raw/instances_train2017.json'\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:fiftyone.utils.coco:Extracting annotations to '/root/fiftyone/coco-2017/raw/instances_train2017.json'\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading 100 images\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:fiftyone.utils.coco:Downloading 100 images\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            " 100% |██████████████████| 100/100 [14.3s elapsed, 0s remaining, 7.9 images/s]      \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:eta.core.utils: 100% |██████████████████| 100/100 [14.3s elapsed, 0s remaining, 7.9 images/s]      \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing annotations for 100 downloaded samples to '/root/fiftyone/coco-2017/train/labels.json'\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:fiftyone.utils.coco:Writing annotations for 100 downloaded samples to '/root/fiftyone/coco-2017/train/labels.json'\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Dataset info written to '/root/fiftyone/coco-2017/info.json'\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:fiftyone.zoo.datasets:Dataset info written to '/root/fiftyone/coco-2017/info.json'\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loading 'coco-2017' split 'train'\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:fiftyone.zoo.datasets:Loading 'coco-2017' split 'train'\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            " 100% |█████████████████| 100/100 [1.7s elapsed, 0s remaining, 59.0 samples/s]         \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:eta.core.utils: 100% |█████████████████| 100/100 [1.7s elapsed, 0s remaining, 59.0 samples/s]         \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Dataset 'coco-2017' created\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:fiftyone.zoo.datasets:Dataset 'coco-2017' created\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "NameError",
          "evalue": "name 'mask_type' is not defined",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-6-144edcde3f85>\u001b[0m in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mdataset\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mCOCODataset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtransform\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mget_transform\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m criterion = MultiBoxLoss(\n\u001b[0m\u001b[1;32m      3\u001b[0m         \u001b[0mnum_classes\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdefault_classes\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m         \u001b[0mpos_threshold\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0.5\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m         \u001b[0mneg_threshold\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0.5\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-1-9f01af96a2be>\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, num_classes, pos_threshold, neg_threshold, negpos_ratio)\u001b[0m\n\u001b[1;32m     38\u001b[0m         \u001b[0muse_class_balanced_conf\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     39\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 40\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmask_type\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmask_type\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"lincomb\"\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     41\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muse_mask_scoring\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     42\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muse_instance_coeff\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mNameError\u001b[0m: name 'mask_type' is not defined"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "next(iter(train_loader))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yu1cFt3h-mCa",
        "outputId": "2fe055f7-5d04-405b-c1a1-6d0f08a9fbc8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "40\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(tensor([[[[ 0.9474,  0.9474,  0.9474,  ...,  0.9988,  0.9988,  1.0159],\n",
              "           [ 0.9474,  0.9474,  0.9474,  ...,  0.9988,  0.9988,  0.9988],\n",
              "           [ 0.9474,  0.9474,  0.9474,  ...,  1.0159,  0.9988,  0.9988],\n",
              "           ...,\n",
              "           [-1.4329, -1.2274, -0.8678,  ...,  1.0673,  1.1187,  1.1187],\n",
              "           [-1.6042, -1.4500, -1.1760,  ...,  1.1529,  1.2728,  1.3584],\n",
              "           [-1.6384, -1.5357, -1.3130,  ...,  1.3070,  1.2899,  1.3070]],\n",
              " \n",
              "          [[ 1.1856,  1.1856,  1.1856,  ...,  1.3081,  1.3081,  1.3256],\n",
              "           [ 1.1856,  1.1856,  1.1856,  ...,  1.3081,  1.3081,  1.3081],\n",
              "           [ 1.1856,  1.1856,  1.1856,  ...,  1.3256,  1.3081,  1.3081],\n",
              "           ...,\n",
              "           [-1.5280, -1.3354, -1.0203,  ...,  1.1155,  1.1681,  1.1681],\n",
              "           [-1.6681, -1.5455, -1.3179,  ...,  1.2031,  1.3256,  1.4132],\n",
              "           [-1.6681, -1.5980, -1.4055,  ...,  1.3606,  1.3431,  1.3606]],\n",
              " \n",
              "          [[ 1.4548,  1.4548,  1.4548,  ...,  1.6291,  1.6465,  1.6640],\n",
              "           [ 1.4548,  1.4548,  1.4548,  ...,  1.6117,  1.6291,  1.6465],\n",
              "           [ 1.4548,  1.4548,  1.4548,  ...,  1.6117,  1.6117,  1.6291],\n",
              "           ...,\n",
              "           [-1.4907, -1.3339, -1.1247,  ...,  0.8797,  0.9668,  0.9668],\n",
              "           [-1.5779, -1.4907, -1.3687,  ...,  0.9668,  1.1237,  1.2108],\n",
              "           [-1.5430, -1.5081, -1.4210,  ...,  1.1237,  1.1411,  1.1585]]],\n",
              " \n",
              " \n",
              "         [[[-1.3987, -1.4500, -1.4672,  ...,  0.2453,  0.2624,  0.2796],\n",
              "           [-1.4500, -1.4672, -1.4672,  ...,  0.2453,  0.2282,  0.2282],\n",
              "           [-1.4672, -1.4500, -1.4672,  ...,  0.2453,  0.2282,  0.2111],\n",
              "           ...,\n",
              "           [ 0.4851,  0.5022,  0.5536,  ...,  0.5878,  0.5536,  0.5364],\n",
              "           [ 0.4508,  0.5022,  0.5536,  ...,  0.5707,  0.5536,  0.5193],\n",
              "           [ 0.5193,  0.5364,  0.5707,  ...,  0.5878,  0.5707,  0.5536]],\n",
              " \n",
              "          [[-1.3179, -1.3354, -1.3354,  ...,  0.9230,  0.9405,  0.9580],\n",
              "           [-1.3179, -1.3354, -1.3179,  ...,  0.9230,  0.9055,  0.9055],\n",
              "           [-1.3354, -1.3179, -1.3179,  ...,  0.9230,  0.9055,  0.8880],\n",
              "           ...,\n",
              "           [ 0.3277,  0.3803,  0.4328,  ...,  0.8704,  0.8354,  0.8179],\n",
              "           [ 0.3102,  0.3627,  0.4153,  ...,  0.8529,  0.8354,  0.8004],\n",
              "           [ 0.3452,  0.3803,  0.4328,  ...,  0.8704,  0.8704,  0.8529]],\n",
              " \n",
              "          [[-1.4559, -1.4733, -1.4907,  ...,  1.0365,  1.0539,  1.0714],\n",
              "           [-1.4733, -1.4907, -1.4559,  ...,  1.0365,  1.0191,  1.0191],\n",
              "           [-1.4733, -1.4559, -1.4559,  ...,  1.0539,  1.0365,  1.0017],\n",
              "           ...,\n",
              "           [-0.4450, -0.4101, -0.3578,  ...,  0.8622,  0.8274,  0.8099],\n",
              "           [-0.5147, -0.4624, -0.4101,  ...,  0.8274,  0.7925,  0.7576],\n",
              "           [-0.5147, -0.4450, -0.3927,  ...,  0.8448,  0.8099,  0.7751]]],\n",
              " \n",
              " \n",
              "         [[[-1.9809, -1.9809, -1.9809,  ..., -1.1075, -1.0562, -0.9020],\n",
              "           [-1.9809, -1.9809, -1.9809,  ...,  0.0569,  0.0569,  0.1083],\n",
              "           [-1.9809, -1.9980, -1.9980,  ...,  0.7933,  0.7248,  0.6734],\n",
              "           ...,\n",
              "           [ 1.1529,  0.8447,  0.5364,  ...,  0.3138,  0.2796,  0.2111],\n",
              "           [ 0.9303,  0.5364,  0.2624,  ...,  0.2111,  0.3309,  0.2624],\n",
              "           [ 0.2111,  0.0056, -0.0629,  ...,  0.2111,  0.3309,  0.2624]],\n",
              " \n",
              "          [[-1.8782, -1.8782, -1.8782,  ..., -1.0028, -0.9678, -0.8277],\n",
              "           [-1.8782, -1.8782, -1.8782,  ...,  0.1702,  0.1702,  0.2227],\n",
              "           [-1.8782, -1.8957, -1.8957,  ...,  0.9230,  0.8529,  0.8004],\n",
              "           ...,\n",
              "           [ 1.2556,  0.9405,  0.5553,  ..., -0.1099, -0.1450, -0.2150],\n",
              "           [ 0.9755,  0.6954,  0.4153,  ..., -0.1800, -0.0574, -0.1275],\n",
              "           [ 0.2227,  0.1527,  0.0651,  ..., -0.1800, -0.0574, -0.1275]],\n",
              " \n",
              "          [[-1.7347, -1.7347, -1.7347,  ..., -0.6715, -0.6018, -0.4101],\n",
              "           [-1.7347, -1.7347, -1.7347,  ...,  0.4962,  0.5136,  0.6008],\n",
              "           [-1.7347, -1.7522, -1.7522,  ...,  1.2108,  1.1585,  1.1237],\n",
              "           ...,\n",
              "           [-0.2532, -0.5147, -0.8458,  ..., -0.3404, -0.3927, -0.4973],\n",
              "           [-0.4450, -0.6541, -0.9678,  ..., -0.3753, -0.2532, -0.3578],\n",
              "           [-1.1421, -1.1596, -1.2816,  ..., -0.3753, -0.2532, -0.3578]]],\n",
              " \n",
              " \n",
              "         ...,\n",
              " \n",
              " \n",
              "         [[[ 0.2282,  0.2282,  0.2453,  ...,  0.0912, -0.3541, -0.2342],\n",
              "           [ 0.1768,  0.1939,  0.2111,  ...,  0.0741, -0.1657, -0.1828],\n",
              "           [ 0.1254,  0.1597,  0.1939,  ...,  0.0569, -0.0801, -0.1486],\n",
              "           ...,\n",
              "           [ 0.8618,  0.5364,  0.2282,  ...,  0.5536,  0.4166,  0.1939],\n",
              "           [ 0.6049,  0.8789,  0.7762,  ...,  0.5707,  0.3309,  0.2796],\n",
              "           [ 1.2043,  1.4269,  1.6667,  ...,  0.2282,  0.1426,  0.2111]],\n",
              " \n",
              "          [[ 0.5203,  0.4853,  0.4503,  ..., -0.1099, -0.4776, -0.0049],\n",
              "           [ 0.4678,  0.4503,  0.4328,  ...,  0.1176, -0.1800, -0.0749],\n",
              "           [ 0.4328,  0.4503,  0.4503,  ...,  0.2052,  0.0126, -0.1099],\n",
              "           ...,\n",
              "           [ 0.6954,  0.6429,  0.4153,  ...,  0.6604,  0.3978,  0.1702],\n",
              "           [ 0.1527,  0.9405,  0.9755,  ...,  0.5903,  0.2052,  0.2052],\n",
              "           [ 1.2381,  1.7808,  1.9734,  ...,  0.1527,  0.0651,  0.2927]],\n",
              " \n",
              "          [[ 1.0017,  0.9842,  0.9842,  ...,  0.3045, -0.0267,  0.4265],\n",
              "           [ 0.9668,  0.9668,  0.9842,  ...,  0.5659,  0.2871,  0.4614],\n",
              "           [ 0.9668,  0.9842,  1.0017,  ...,  0.6182,  0.4091,  0.3916],\n",
              "           ...,\n",
              "           [ 1.1237,  0.9145,  0.4265,  ...,  0.8971,  0.6356,  0.4439],\n",
              "           [ 0.6705,  1.2457,  1.1062,  ...,  0.7576,  0.4091,  0.4091],\n",
              "           [ 1.5942,  1.9603,  2.0823,  ...,  0.3045,  0.1999,  0.4265]]],\n",
              " \n",
              " \n",
              "         [[[ 1.5639,  1.5639,  1.5810,  ...,  1.2214,  1.2043,  1.1872],\n",
              "           [ 1.5639,  1.5639,  1.5810,  ...,  1.2214,  1.2214,  1.1872],\n",
              "           [ 1.5639,  1.5639,  1.5810,  ...,  1.2557,  1.1872,  1.1015],\n",
              "           ...,\n",
              "           [-1.3644, -1.3302, -1.2788,  ..., -0.6794, -0.7308, -0.5938],\n",
              "           [-1.3473, -1.3644, -1.3302,  ..., -0.6623, -0.6623, -0.5424],\n",
              "           [-1.3644, -1.3815, -1.3473,  ..., -0.7308, -0.7308, -0.6794]],\n",
              " \n",
              "          [[ 1.9384,  1.9384,  1.9559,  ...,  1.6758,  1.6583,  1.6408],\n",
              "           [ 1.9384,  1.9384,  1.9559,  ...,  1.6758,  1.6758,  1.6408],\n",
              "           [ 1.9384,  1.9384,  1.9559,  ...,  1.7108,  1.6408,  1.5532],\n",
              "           ...,\n",
              "           [-0.6527, -0.6527, -0.6176,  ..., -0.0399, -0.1099,  0.0126],\n",
              "           [-0.7052, -0.7227, -0.6877,  ..., -0.0049, -0.0224,  0.0651],\n",
              "           [-0.7227, -0.7402, -0.7052,  ..., -0.0749, -0.0924, -0.0749]],\n",
              " \n",
              "          [[ 2.4308,  2.4308,  2.4483,  ...,  2.0648,  2.0474,  2.0300],\n",
              "           [ 2.4308,  2.4308,  2.4483,  ...,  2.0648,  2.0648,  2.0300],\n",
              "           [ 2.4308,  2.4308,  2.4483,  ...,  2.0997,  2.0300,  1.9428],\n",
              "           ...,\n",
              "           [ 0.1476,  0.1302,  0.1476,  ...,  0.5659,  0.4788,  0.6182],\n",
              "           [ 0.0256,  0.0082,  0.0256,  ...,  0.5834,  0.5659,  0.6705],\n",
              "           [ 0.0082, -0.0092,  0.0082,  ...,  0.5136,  0.5136,  0.5311]]],\n",
              " \n",
              " \n",
              "         [[[-1.9124, -1.9295, -1.9638,  ..., -1.6727, -1.6727, -1.6727],\n",
              "           [-1.9295, -1.9467, -1.9809,  ..., -1.6042, -1.5870, -1.5870],\n",
              "           [-1.9467, -1.9638, -1.9809,  ..., -1.5185, -1.5014, -1.5014],\n",
              "           ...,\n",
              "           [-1.4843, -1.4843, -1.4672,  ...,  1.5125,  1.5810,  1.6324],\n",
              "           [-1.4500, -1.4500, -1.4500,  ...,  1.5468,  1.5982,  1.6495],\n",
              "           [-1.4158, -1.4158, -1.4158,  ...,  1.6153,  1.6667,  1.7180]],\n",
              " \n",
              "          [[-1.7556, -1.7731, -1.8081,  ..., -1.5980, -1.5980, -1.5980],\n",
              "           [-1.7731, -1.7906, -1.8256,  ..., -1.5105, -1.5105, -1.5105],\n",
              "           [-1.7906, -1.8081, -1.8256,  ..., -1.4055, -1.3880, -1.3880],\n",
              "           ...,\n",
              "           [-1.0553, -1.0553, -1.0378,  ...,  1.3606,  1.4307,  1.4832],\n",
              "           [-1.0203, -1.0203, -1.0203,  ...,  1.3957,  1.4482,  1.5007],\n",
              "           [-1.0028, -1.0028, -0.9853,  ...,  1.4657,  1.5182,  1.5707]],\n",
              " \n",
              "          [[-1.5081, -1.5256, -1.5604,  ..., -1.4036, -1.4036, -1.4036],\n",
              "           [-1.5256, -1.5430, -1.5779,  ..., -1.3861, -1.3687, -1.3687],\n",
              "           [-1.5430, -1.5604, -1.5779,  ..., -1.3339, -1.3164, -1.3164],\n",
              "           ...,\n",
              "           [ 0.8274,  0.8274,  0.8448,  ...,  1.3328,  1.4025,  1.4548],\n",
              "           [ 0.8797,  0.8797,  0.8622,  ...,  1.3677,  1.4200,  1.4722],\n",
              "           [ 0.9319,  0.9145,  0.8971,  ...,  1.4374,  1.4897,  1.5420]]]]),\n",
              " {'bboxes': tensor([[[ 16., 114., 281., 273.,  28.],\n",
              "           [177., 140., 368., 247.,  28.],\n",
              "           [253., 133., 413., 218.,  28.],\n",
              "           ...,\n",
              "           [145., 392., 224., 466.,  31.],\n",
              "           [ 56., 202., 714., 347.,   1.],\n",
              "           [480., 130., 610., 205.,  28.]],\n",
              "  \n",
              "          [[  1., 116., 799., 507.,  81.],\n",
              "           [  0., 160., 380., 277.,  81.],\n",
              "           [197., 181., 436., 275.,  90.],\n",
              "           ...,\n",
              "           [  0.,   0.,   0.,   0.,   0.],\n",
              "           [  0.,   0.,   0.,   0.,   0.],\n",
              "           [  0.,   0.,   0.,   0.,   0.]],\n",
              "  \n",
              "          [[481.,  72., 749., 429.,  25.],\n",
              "           [ 66., 428., 231., 494.,  25.],\n",
              "           [  0.,   0.,   0.,   0.,   0.],\n",
              "           ...,\n",
              "           [  0.,   0.,   0.,   0.,   0.],\n",
              "           [  0.,   0.,   0.,   0.,   0.],\n",
              "           [  0.,   0.,   0.,   0.,   0.]],\n",
              "  \n",
              "          ...,\n",
              "  \n",
              "          [[283., 105., 624., 414.,  23.],\n",
              "           [  0.,   0.,   0.,   0.,   0.],\n",
              "           [  0.,   0.,   0.,   0.,   0.],\n",
              "           ...,\n",
              "           [  0.,   0.,   0.,   0.,   0.],\n",
              "           [  0.,   0.,   0.,   0.,   0.],\n",
              "           [  0.,   0.,   0.,   0.,   0.]],\n",
              "  \n",
              "          [[168.,  34., 335., 509.,   1.],\n",
              "           [  0.,  56., 231., 511.,   1.],\n",
              "           [ 57.,  18., 260., 337.,   1.],\n",
              "           ...,\n",
              "           [  0.,   0.,   0.,   0.,   0.],\n",
              "           [  0.,   0.,   0.,   0.,   0.],\n",
              "           [  0.,   0.,   0.,   0.,   0.]],\n",
              "  \n",
              "          [[  1., 157.,  86., 509.,   1.],\n",
              "           [  0.,  83., 110., 511.,   1.],\n",
              "           [131., 159., 428., 504.,  25.],\n",
              "           ...,\n",
              "           [  0.,   0.,   0.,   0.,   0.],\n",
              "           [  0.,   0.,   0.,   0.,   0.],\n",
              "           [  0.,   0.,   0.,   0.,   0.]]]),\n",
              "  'masks': tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           ...,\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
              "  \n",
              "  \n",
              "          [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           ...,\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
              "  \n",
              "  \n",
              "          [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           ...,\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
              "  \n",
              "  \n",
              "          ...,\n",
              "  \n",
              "  \n",
              "          [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           ...,\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
              "  \n",
              "  \n",
              "          [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           ...,\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
              "  \n",
              "  \n",
              "          [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           ...,\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]],\n",
              "  \n",
              "           [[0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            ...,\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.],\n",
              "            [0., 0., 0.,  ..., 0., 0., 0.]]]]),\n",
              "  'num_crowds': tensor([0])})"
            ]
          },
          "metadata": {},
          "execution_count": 77
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def train():\n",
        "    dataset = COCODataset(transform=get_transform())\n",
        "\n",
        "    train_loader = torch.utils.data.DataLoader(\n",
        "        dataset, batch_size=10, shuffle=True, collate_fn=collate_fn\n",
        "    )\n",
        "\n",
        "    yolact = Yolact()\n",
        "    yolact.init_backbone_weights()\n",
        "\n",
        "    optim = torch.optim.SGD(yolact.parameters(), lr=0.001, momentum=0.9,\n",
        "                            weight_decay=0.0005)\n",
        "\n",
        "    criterion = MultiBoxLoss(\n",
        "        num_classes=len(dataset.dataset.default_classes),\n",
        "        pos_threshold=0.5,\n",
        "        neg_threshold=0.5,\n",
        "        # When using ohem, the ratio between positives and negatives\n",
        "        # (3 means 3 negatives to 1 positive)\n",
        "        negpos_ratio=3,\n",
        "    )\n",
        "\n",
        "    # freeze bn\n",
        "\n",
        "    num_epochs = 100\n",
        "\n",
        "\n",
        "    for epoch in range(num_epochs):\n",
        "        yolact.train()\n",
        "\n",
        "        for batch_idx, (images, targets) in enumerate(train_loader):\n",
        "            images, targets = images.cuda(), targets.cuda()\n",
        "            bboxes, masks, num_crowds = targets[\"bboxes\"], targets[\"masks\"], targets[\"num_crowds\"]\n",
        "\n",
        "            output = yolact(images)\n",
        "            loss = criterion(output, targets, masks, num_crowds)\n",
        "\n",
        "            optim.zero_grad()\n",
        "            loss.backward()\n",
        "            optim.step()\n",
        "\n",
        "            if batch_idx % 10 == 0:\n",
        "                print(f\"epoch: {epoch}, batch: {batch_idx}, loss: {loss.item()}\")"
      ],
      "metadata": {
        "id": "GvLVpOq_LJ-x"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "HHL1h7r-bZAM"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "V100",
      "machine_shape": "hm",
      "provenance": [],
      "authorship_tag": "ABX9TyM+xgqPgXfwDs2iljGZyN7X",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}