{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "\n",
    "import extractors\n",
    "\n",
    "\n",
    "class PSPModule(nn.Module):\n",
    "    def __init__(self, features, out_features=1024, sizes=(1, 2, 3, 6)):\n",
    "        super().__init__()\n",
    "        self.stages = []\n",
    "        self.stages = nn.ModuleList([self._make_stage(features, size) for size in sizes])\n",
    "        self.bottleneck = nn.Conv2d(features * (len(sizes) + 1), out_features, kernel_size=1)\n",
    "        self.relu = nn.ReLU()\n",
    "\n",
    "    def _make_stage(self, features, size):\n",
    "        prior = nn.AdaptiveAvgPool2d(output_size=(size, size))\n",
    "        conv = nn.Conv2d(features, features, kernel_size=1, bias=False)\n",
    "        return nn.Sequential(prior, conv)\n",
    "\n",
    "    def forward(self, feats):\n",
    "        h, w = feats.size(2), feats.size(3)\n",
    "        priors = [F.upsample(input=stage(feats), size=(h, w), mode='bilinear') for stage in self.stages] + [feats]\n",
    "        bottle = self.bottleneck(torch.cat(priors, 1))\n",
    "        return self.relu(bottle)\n",
    "\n",
    "\n",
    "class PSPUpsample(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super().__init__()\n",
    "        self.conv = nn.Sequential(\n",
    "            nn.Conv2d(in_channels, out_channels, 3, padding=1),\n",
    "            nn.BatchNorm2d(out_channels),\n",
    "            nn.PReLU()\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        h, w = 2 * x.size(2), 2 * x.size(3)\n",
    "        p = F.upsample(input=x, size=(h, w), mode='bilinear')\n",
    "        return self.conv(p)\n",
    "\n",
    "\n",
    "class PSPNet(nn.Module):\n",
    "    def __init__(self, n_classes=18, sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet34',\n",
    "                 pretrained=True):\n",
    "        super().__init__()\n",
    "        self.feats = getattr(extractors, backend)(pretrained)\n",
    "        self.psp = PSPModule(psp_size, 1024, sizes)\n",
    "        self.drop_1 = nn.Dropout2d(p=0.3)\n",
    "\n",
    "        self.up_1 = PSPUpsample(1024, 256)\n",
    "        self.up_2 = PSPUpsample(256, 64)\n",
    "        self.up_3 = PSPUpsample(64, 64)\n",
    "\n",
    "        self.drop_2 = nn.Dropout2d(p=0.15)\n",
    "        self.final = nn.Sequential(\n",
    "            nn.Conv2d(64, n_classes, kernel_size=1),\n",
    "            nn.LogSoftmax()\n",
    "        )\n",
    "\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Linear(deep_features_size, 256),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(256, n_classes)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        f, class_f = self.feats(x) \n",
    "        p = self.psp(f)\n",
    "        p = self.drop_1(p)\n",
    "\n",
    "        p = self.up_1(p)\n",
    "        p = self.drop_2(p)\n",
    "\n",
    "        p = self.up_2(p)\n",
    "        p = self.drop_2(p)\n",
    "\n",
    "        p = self.up_3(p)\n",
    "        p = self.drop_2(p)\n",
    "\n",
    "        auxiliary = F.adaptive_max_pool2d(input=class_f, output_size=(1, 1)).view(-1, class_f.size(1))\n",
    "\n",
    "        return self.final(p), self.classifier(auxiliary)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
