{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "\n",
    "from incendio.core import BaseModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(BaseModel):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.pad = nn.ReflectionPad2d(2)\n",
    "        self.conv1 = nn.Conv2d(3, 8, kernel_size=5)\n",
    "        self.conv2 = nn.Conv2d(8, 16, kernel_size=5)\n",
    "        self.adapt = nn.AdaptiveAvgPool2d((1, 1))\n",
    "        self.fc = nn.Linear(16, 1)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.conv1(self.pad(x))\n",
    "        x = F.leaky_relu(x)\n",
    "        x = self.conv2(self.pad(x))\n",
    "        x = F.leaky_relu(x)\n",
    "        x = self.adapt(x)\n",
    "        x = self.fc(x.squeeze())\n",
    "        return torch.sigmoid(x)\n",
    "    \n",
    "    \n",
    "class Model2(BaseModel):\n",
    "\n",
    "    def __init__(self, c_in, c_outs):\n",
    "        super().__init__()\n",
    "        dims = [c_in] + c_outs\n",
    "        self.enc = nn.Sequential(*[nn.Sequential(nn.ReflectionPad2d(2),\n",
    "                                   nn.Conv2d(*(c_in, c_out), kernel_size=5),\n",
    "                                   nn.LeakyReLU())\n",
    "                                   for c_in, c_out in zip(dims, dims[1:])])\n",
    "        \n",
    "        self.pool = nn.AdaptiveAvgPool2d((1, 1))\n",
    "        self.fc = nn.Linear(c_outs[-1], 1)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.enc(x)\n",
    "        x = self.pool(x)\n",
    "        print(x.shape)\n",
    "        x = x.squeeze()\n",
    "        print(x.shape)\n",
    "        return torch.sigmoid(self.fc(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Model2(\n",
       "  (enc): Sequential(\n",
       "    (0): Sequential(\n",
       "      (0): ReflectionPad2d((2, 2, 2, 2))\n",
       "      (1): Conv2d(3, 8, kernel_size=(5, 5), stride=(1, 1))\n",
       "      (2): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (1): Sequential(\n",
       "      (0): ReflectionPad2d((2, 2, 2, 2))\n",
       "      (1): Conv2d(8, 16, kernel_size=(5, 5), stride=(1, 1))\n",
       "      (2): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (2): Sequential(\n",
       "      (0): ReflectionPad2d((2, 2, 2, 2))\n",
       "      (1): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1))\n",
       "      (2): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "  )\n",
       "  (pool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
       "  (fc): Linear(in_features=32, out_features=1, bias=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m2 = Model2(3, [8, 16, 32])\n",
    "m2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Model(\n",
       "  (pad): ReflectionPad2d((2, 2, 2, 2))\n",
       "  (conv1): Conv2d(3, 8, kernel_size=(5, 5), stride=(1, 1))\n",
       "  (conv2): Conv2d(8, 16, kernel_size=(5, 5), stride=(1, 1))\n",
       "  (adapt): AdaptiveAvgPool2d(output_size=(1, 1))\n",
       "  (fc): Linear(in_features=16, out_features=1, bias=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m1 = Model()\n",
    "m1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[  6., 210., 214.,  38.],\n",
       "          [176., 197.,  80., 112.],\n",
       "          [168.,  81., 106., 252.],\n",
       "          [202.,  38., 147., 102.]],\n",
       "\n",
       "         [[ 28.,  64., 112., 143.],\n",
       "          [152.,  87., 219., 151.],\n",
       "          [123., 154.,  79.,  27.],\n",
       "          [127., 198., 100., 223.]],\n",
       "\n",
       "         [[253., 150., 155., 243.],\n",
       "          [222.,  96., 219.,  40.],\n",
       "          [ 75., 251., 234., 149.],\n",
       "          [112., 234., 204.,  13.]]],\n",
       "\n",
       "\n",
       "        [[[248., 236., 113., 144.],\n",
       "          [223.,  34., 192., 156.],\n",
       "          [248., 183.,  35.,  35.],\n",
       "          [101.,  64.,  25., 203.]],\n",
       "\n",
       "         [[ 23., 145.,  52., 237.],\n",
       "          [119., 254., 116., 180.],\n",
       "          [ 70., 199., 149.,   5.],\n",
       "          [ 78.,  58.,  22.,  34.]],\n",
       "\n",
       "         [[221., 226.,  48., 109.],\n",
       "          [212., 146.,  67.,   0.],\n",
       "          [ 65.,  45., 175.,  85.],\n",
       "          [233., 187., 190.,  22.]]]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randint(255, (2, 3, 4, 4)).float()\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 4, 4])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_img(img):\n",
    "    plt.imshow(img.permute(1, 2, 0) / 255)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 4712 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_img(x[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 4, 4])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 3, 4, 4])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 3, 8, 8])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pad = nn.ReflectionPad2d(2)\n",
    "x_pad = pad(x)\n",
    "x_pad.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 4372 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_img(x_pad[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReflectionPaddedConv2d(nn.Module):\n",
    "    \n",
    "    def __init__(self, in_channels, out_channels, padding=1, \n",
    "                 kernel_size=3, stride=1, bias=True, **kwargs):\n",
    "        super().__init__()\n",
    "        self.reflect = nn.ReflectionPad2d(padding)\n",
    "        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size,\n",
    "                              stride, padding=0, bias=bias)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.reflect(x)\n",
    "        return self.conv(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([0., 0., 0.], requires_grad=True)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rc = ReflectionPaddedConv2d(in_channels=3, out_channels=3)\n",
    "nn.init.constant_(rc.conv.weight, 1)\n",
    "nn.init.constant_(rc.conv.bias, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 3, 4, 4])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "r = nn.Conv2d(3, 3, kernel_size=3, padding=1, padding_mode='zeros', bias=True)\n",
    "nn.init.constant_(r.weight, 1)\n",
    "nn.init.constant_(r.bias, 0)\n",
    "x_p = r(x)\n",
    "x_p.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Applies a 2D convolution over an input signal composed of several input\n",
      "    planes.\n",
      "\n",
      "    In the simplest case, the output value of the layer with input size\n",
      "    :math:`(N, C_{\\text{in}}, H, W)` and output :math:`(N, C_{\\text{out}}, H_{\\text{out}}, W_{\\text{out}})`\n",
      "    can be precisely described as:\n",
      "\n",
      "    .. math::\n",
      "        \\text{out}(N_i, C_{\\text{out}_j}) = \\text{bias}(C_{\\text{out}_j}) +\n",
      "        \\sum_{k = 0}^{C_{\\text{in}} - 1} \\text{weight}(C_{\\text{out}_j}, k) \\star \\text{input}(N_i, k)\n",
      "\n",
      "\n",
      "    where :math:`\\star` is the valid 2D `cross-correlation`_ operator,\n",
      "    :math:`N` is a batch size, :math:`C` denotes a number of channels,\n",
      "    :math:`H` is a height of input planes in pixels, and :math:`W` is\n",
      "    width in pixels.\n",
      "\n",
      "    * :attr:`stride` controls the stride for the cross-correlation, a single\n",
      "      number or a tuple.\n",
      "\n",
      "    * :attr:`padding` controls the amount of implicit zero-paddings on both\n",
      "      sides for :attr:`padding` number of points for each dimension.\n",
      "\n",
      "    * :attr:`dilation` controls the spacing between the kernel points; also\n",
      "      known as the à trous algorithm. It is harder to describe, but this `link`_\n",
      "      has a nice visualization of what :attr:`dilation` does.\n",
      "\n",
      "    * :attr:`groups` controls the connections between inputs and outputs.\n",
      "      :attr:`in_channels` and :attr:`out_channels` must both be divisible by\n",
      "      :attr:`groups`. For example,\n",
      "\n",
      "        * At groups=1, all inputs are convolved to all outputs.\n",
      "        * At groups=2, the operation becomes equivalent to having two conv\n",
      "          layers side by side, each seeing half the input channels,\n",
      "          and producing half the output channels, and both subsequently\n",
      "          concatenated.\n",
      "        * At groups= :attr:`in_channels`, each input channel is convolved with\n",
      "          its own set of filters, of size:\n",
      "          :math:`\\left\\lfloor\\frac{out\\_channels}{in\\_channels}\\right\\rfloor`.\n",
      "\n",
      "    The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, :attr:`dilation` can either be:\n",
      "\n",
      "        - a single ``int`` -- in which case the same value is used for the height and width dimension\n",
      "        - a ``tuple`` of two ints -- in which case, the first `int` is used for the height dimension,\n",
      "          and the second `int` for the width dimension\n",
      "\n",
      "    .. note::\n",
      "\n",
      "         Depending of the size of your kernel, several (of the last)\n",
      "         columns of the input might be lost, because it is a valid `cross-correlation`_,\n",
      "         and not a full `cross-correlation`_.\n",
      "         It is up to the user to add proper padding.\n",
      "\n",
      "    .. note::\n",
      "\n",
      "        When `groups == in_channels` and `out_channels == K * in_channels`,\n",
      "        where `K` is a positive integer, this operation is also termed in\n",
      "        literature as depthwise convolution.\n",
      "\n",
      "        In other words, for an input of size :math:`(N, C_{in}, H_{in}, W_{in})`,\n",
      "        a depthwise convolution with a depthwise multiplier `K`, can be constructed by arguments\n",
      "        :math:`(in\\_channels=C_{in}, out\\_channels=C_{in} \\times K, ..., groups=C_{in})`.\n",
      "\n",
      "    .. include:: cudnn_deterministic.rst\n",
      "\n",
      "    Args:\n",
      "        in_channels (int): Number of channels in the input image\n",
      "        out_channels (int): Number of channels produced by the convolution\n",
      "        kernel_size (int or tuple): Size of the convolving kernel\n",
      "        stride (int or tuple, optional): Stride of the convolution. Default: 1\n",
      "        padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0\n",
      "        padding_mode (string, optional). Accepted values `zeros` and `circular` Default: `zeros`\n",
      "        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1\n",
      "        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1\n",
      "        bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``\n",
      "\n",
      "    Shape:\n",
      "        - Input: :math:`(N, C_{in}, H_{in}, W_{in})`\n",
      "        - Output: :math:`(N, C_{out}, H_{out}, W_{out})` where\n",
      "\n",
      "          .. math::\n",
      "              H_{out} = \\left\\lfloor\\frac{H_{in}  + 2 \\times \\text{padding}[0] - \\text{dilation}[0]\n",
      "                        \\times (\\text{kernel\\_size}[0] - 1) - 1}{\\text{stride}[0]} + 1\\right\\rfloor\n",
      "\n",
      "          .. math::\n",
      "              W_{out} = \\left\\lfloor\\frac{W_{in}  + 2 \\times \\text{padding}[1] - \\text{dilation}[1]\n",
      "                        \\times (\\text{kernel\\_size}[1] - 1) - 1}{\\text{stride}[1]} + 1\\right\\rfloor\n",
      "\n",
      "    Attributes:\n",
      "        weight (Tensor): the learnable weights of the module of shape\n",
      "                         :math:`(\\text{out\\_channels}, \\frac{\\text{in\\_channels}}{\\text{groups}},`\n",
      "                         :math:`\\text{kernel\\_size[0]}, \\text{kernel\\_size[1]})`.\n",
      "                         The values of these weights are sampled from\n",
      "                         :math:`\\mathcal{U}(-\\sqrt{k}, \\sqrt{k})` where\n",
      "                         :math:`k = \\frac{1}{C_\\text{in} * \\prod_{i=0}^{1}\\text{kernel\\_size}[i]}`\n",
      "        bias (Tensor):   the learnable bias of the module of shape (out_channels). If :attr:`bias` is ``True``,\n",
      "                         then the values of these weights are\n",
      "                         sampled from :math:`\\mathcal{U}(-\\sqrt{k}, \\sqrt{k})` where\n",
      "                         :math:`k = \\frac{1}{C_\\text{in} * \\prod_{i=0}^{1}\\text{kernel\\_size}[i]}`\n",
      "\n",
      "    Examples::\n",
      "\n",
      "        >>> # With square kernels and equal stride\n",
      "        >>> m = nn.Conv2d(16, 33, 3, stride=2)\n",
      "        >>> # non-square kernels and unequal stride and with padding\n",
      "        >>> m = nn.Conv2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2))\n",
      "        >>> # non-square kernels and unequal stride and with padding and dilation\n",
      "        >>> m = nn.Conv2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2), dilation=(3, 1))\n",
      "        >>> input = torch.randn(20, 16, 50, 100)\n",
      "        >>> output = m(input)\n",
      "\n",
      "    .. _cross-correlation:\n",
      "        https://en.wikipedia.org/wiki/Cross-correlation\n",
      "\n",
      "    .. _link:\n",
      "        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md\n",
      "    \n"
     ]
    }
   ],
   "source": [
    "print(nn.Conv2d.__doc__)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
