Resnet_UNet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "constant-newport",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision import models\n",
    "import torch.nn as nn\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "blank-display",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResNetUNet(nn.Module):\n",
    "    def __init__(self, n_class):\n",
    "        super().__init__()\n",
    "\n",
    "        self.base_model = models.resnet18(pretrained=True)\n",
    "        for param in self.base_model.parameters():\n",
    "            param.requires_grad = False\n",
    "        self.base_layers = list(self.base_model.children())\n",
    "\n",
    "        self.layer0 = nn.Sequential(*self.base_layers[:3]) # size=(N, 64, x.H/2, x.W/2)\n",
    "        self.layer0_1x1 = self.convrelu(64, 64, 1, 0)\n",
    "        self.layer1 = nn.Sequential(*self.base_layers[3:5]) # size=(N, 64, x.H/4, x.W/4)\n",
    "        self.layer1_1x1 = self.convrelu(64, 64, 1, 0)\n",
    "        self.layer2 = self.base_layers[5]  # size=(N, 128, x.H/8, x.W/8)\n",
    "        self.layer2_1x1 = self.convrelu(128, 128, 1, 0)\n",
    "        self.layer3 = self.base_layers[6]  # size=(N, 256, x.H/16, x.W/16)\n",
    "        self.layer3_1x1 = self.convrelu(256, 256, 1, 0)\n",
    "        self.layer4 = self.base_layers[7]  # size=(N, 512, x.H/32, x.W/32)\n",
    "        self.layer4_1x1 = self.convrelu(512, 512, 1, 0)\n",
    "\n",
    "        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n",
    "\n",
    "        self.conv_up3 = self.convrelu(256 + 512, 512, 3, 1)\n",
    "        self.conv_up2 = self.convrelu(128 + 512, 256, 3, 1)\n",
    "        self.conv_up1 = self.convrelu(64 + 256, 256, 3, 1)\n",
    "        self.conv_up0 = self.convrelu(64 + 256, 128, 3, 1)\n",
    "\n",
    "        self.conv_original_size0 = self.convrelu(3, 64, 3, 1)\n",
    "        self.conv_original_size1 = self.convrelu(64, 64, 3, 1)\n",
    "        self.conv_original_size2 = self.convrelu(64 + 128, 64, 3, 1)\n",
    "\n",
    "        self.conv_last = nn.Conv2d(64, n_class, 1)\n",
    "\n",
    "    def forward(self, input):\n",
    "        x_original = self.conv_original_size0(input)\n",
    "        x_original = self.conv_original_size1(x_original)\n",
    "\n",
    "        layer0 = self.layer0(input)\n",
    "        layer1 = self.layer1(layer0)\n",
    "        layer2 = self.layer2(layer1)\n",
    "        layer3 = self.layer3(layer2)\n",
    "        layer4 = self.layer4(layer3)\n",
    "\n",
    "        layer4 = self.layer4_1x1(layer4)\n",
    "        x = self.upsample(layer4)\n",
    "        layer3 = self.layer3_1x1(layer3)\n",
    "        x = torch.cat([x, layer3], dim=1)\n",
    "        x = self.conv_up3(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        layer2 = self.layer2_1x1(layer2)\n",
    "        x = torch.cat([x, layer2], dim=1)\n",
    "        x = self.conv_up2(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        layer1 = self.layer1_1x1(layer1)\n",
    "        x = torch.cat([x, layer1], dim=1)\n",
    "        x = self.conv_up1(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        layer0 = self.layer0_1x1(layer0)\n",
    "        x = torch.cat([x, layer0], dim=1)\n",
    "        x = self.conv_up0(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        x = torch.cat([x, x_original], dim=1)\n",
    "        x = self.conv_original_size2(x)\n",
    "\n",
    "        out = self.conv_last(x)\n",
    "\n",
    "        return out\n",
    "    \n",
    "    def convrelu(self,in_channels, out_channels, kernel, padding):\n",
    "        return nn.Sequential(\n",
    "            nn.Conv2d(in_channels, out_channels, kernel, padding=padding),\n",
    "            nn.ReLU(inplace=True),\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aerial-south",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}