{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp layers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Layers\n",
    "\n",
    "> Custom activations, layers, and layer blocks are contained in this module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Used in notebook but not needed in package.\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Activations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class GRelu(nn.Module):\n",
    "    \"\"\"Generic ReLU.\"\"\"\n",
    "\n",
    "    def __init__(self, leak=0.0, max=float('inf'), sub=0.0):\n",
    "        super().__init__()\n",
    "        self.leak = leak\n",
    "        self.max = max\n",
    "        self.sub = sub\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"Check which operations are necessary to save computation.\"\"\"\n",
    "        x = F.leaky_relu(x, self.leak) if self.leak else F.relu(x)\n",
    "        if self.sub:\n",
    "            x -= self.sub\n",
    "        if self.max:\n",
    "            x = torch.clamp(x, max=self.max)\n",
    "        return x\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f'GReLU(leak={self.leak}, max={self.max}, sub={self.sub})'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "JRelu = GRelu(leak=.1, sub=.4, max=6.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class Mish(nn.Module):\n",
    "    \"\"\"OOP form of mish activation.\n",
    "    \n",
    "    Mish: A Self Regularized Non-Monotonic Neural Activation Function\n",
    "    https://arxiv.org/pdf/1908.08681v1.pdf\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return x * torch.tanh(F.softplus(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def mish(x):\n",
    "    \"\"\"Functional form of mish activation.\n",
    "    \n",
    "    Mish: A Self Regularized Non-Monotonic Neural Activation Function\n",
    "    https://arxiv.org/pdf/1908.08681v1.pdf\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x: torch.Tensor[float]\n",
    "        Input tensor.\n",
    "    Returns\n",
    "    -------\n",
    "    torch.Tensor[float]: Tensor of same shape as input x.\n",
    "    \"\"\"\n",
    "    return x * torch.tanh(F.softplus(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_activations(z, a, mode='scatter', **kwargs):\n",
    "    \"\"\"Plot an input tensor and its corresponding activations.  Both tensors\n",
    "    will be flattened for plotting.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    z: tf.Tensor\n",
    "        Tensor containing values to plot on the x axis (we can often think of\n",
    "        this as the output of a linear layer, where z=f(x) and a=mish(z)).\n",
    "    a: tf.Tensor\n",
    "        Tensor containing values to plot on y axis.\n",
    "    mode: str\n",
    "        'scatter' for scatter plot or 'plot' for line plot.\n",
    "    kwargs: Values to be passed to the matplotlib plotting function, such as \n",
    "        's' when in 'scatter' mode or 'lw' in 'plot' mode.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    None\n",
    "    \"\"\"\n",
    "    plt_func = getattr(plt, mode)\n",
    "    kwargs = kwargs or {}\n",
    "    if mode == 'scatter' and not kwargs:\n",
    "        kwargs = {'s': .75}\n",
    "    plt_func(z.numpy().flatten(), a.numpy().flatten(), **kwargs)\n",
    "    plt.axvline(0, lw=.5, alpha=.5)\n",
    "    plt.axhline(0, lw=.5, alpha=.5)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.arange(-5, 5, .05)\n",
    "a = mish(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 10384 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_activations(x, a, 'plot')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Layer Blocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class ConvBlock(nn.Module):\n",
    "    \"\"\"Create a convolutional block optionally followed by a batch norm layer.\n",
    "    \"\"\"\n",
    "        \n",
    "    def __init__(self, c_in, c_out, kernel_size=3, norm=True, activation=JRelu, \n",
    "                 **kwargs):\n",
    "        \"\"\"        \n",
    "        Parameters\n",
    "        -----------\n",
    "        c_in: int\n",
    "            # of input channels.\n",
    "        c_out: int\n",
    "            # of output channels.\n",
    "        kernel_size: int\n",
    "            Size of kernel in conv2d layer. An integer argument will be used \n",
    "            as both the height and width.\n",
    "        norm: bool\n",
    "            If True, include a batch norm layer after the conv layer. If False,\n",
    "            no norm layer will be used. Note that batch norm has learnable \n",
    "            affine parameters which remove the need for a bias in the preceding\n",
    "            conv layer. When batch norm is not used, however, the conv layer \n",
    "            will include a bias term.\n",
    "        activation: nn.Module\n",
    "            Activation function to use at the end of the convolutional block.\n",
    "            (In some cases such as our ResBlock implementation, we pass in None\n",
    "            so that an extra addition can be performed before the final\n",
    "            activation.) Do not use the functional form here as it will be\n",
    "            added to a sequential object.\n",
    "        kwargs: any\n",
    "            Additional keyword args are passed to Conv2d. Useful kwargs include \n",
    "            stride, and padding (see pytorch docs for nn.Conv2d).\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.norm = norm\n",
    "        layers = [nn.Conv2d(c_in, c_out, kernel_size, bias=not norm, **kwargs)]\n",
    "        if norm: \n",
    "            layers.append(nn.BatchNorm2d(c_out))\n",
    "        if activation is not None: \n",
    "            layers.append(activation)\n",
    "        self.block = nn.Sequential(*layers)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        self.block(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ConvBlock(\n",
       "  (block): Sequential(\n",
       "    (0): Conv2d(3, 5, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (1): GReLU(leak=0.1, max=6.0, sub=0.4)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ConvBlock(3, 5, norm=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class ResBlock(nn.Module):\n",
    "\n",
    "    def __init__(self, c_in, activation=JRelu, f=3, stride=1, pad=1,\n",
    "                 skip_size=2, norm=True):\n",
    "        \"\"\"Residual block to be used in CycleGenerator. Note that f, stride,\n",
    "        and pad must be selected such that the height and width of the input\n",
    "        remain the same.\n",
    "\n",
    "        Parameters\n",
    "        -----------\n",
    "        c_in: int\n",
    "            # of input channels.\n",
    "        activation: callable\n",
    "            Activation function to use.\n",
    "        f: int\n",
    "            Size of filter (f x f) used in convolution. Default 3.\n",
    "        stride: int\n",
    "            # of pixels the filter moves between each convolution. Default 1.\n",
    "        pad: int\n",
    "            Pixel padding around the input. Default 1.\n",
    "        skip_size: int\n",
    "            Number of conv blocks inside the skip connection (default 2).\n",
    "            ResNet paper notes that skipping a single layer did not show\n",
    "            noticeable improvements.\n",
    "        norm: bool\n",
    "            Specifies whether to include a batch norm layer after each conv\n",
    "            layer.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.skip_size = skip_size\n",
    "        self.layers = nn.ModuleList([\n",
    "            ConvBlock(c_in, c_in, norm=norm, activation=None, kernel_size=f,\n",
    "                      stride=stride, padding=pad)\n",
    "            for i in range(skip_size)\n",
    "        ])\n",
    "        self.activation = activation\n",
    "\n",
    "    def forward(self, x):\n",
    "        x_out = x\n",
    "        for i, layer in enumerate(self.layers):\n",
    "            x_out = layer(x_out)\n",
    "\n",
    "            # Final activation must be applied after addition.\n",
    "            if i != self.skip_size - 1:\n",
    "                x_out = self.activation(x_out)\n",
    "\n",
    "        return self.activation(x + x_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResBlock(\n",
       "  (layers): ModuleList(\n",
       "    (0): ConvBlock(\n",
       "      (block): Sequential(\n",
       "        (0): Conv2d(4, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): BatchNorm2d(4, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): ConvBlock(\n",
       "      (block): Sequential(\n",
       "        (0): Conv2d(4, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): BatchNorm2d(4, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (activation): GReLU(leak=0.1, max=6.0, sub=0.4)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ResBlock(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResBlock(\n",
       "  (layers): ModuleList(\n",
       "    (0): ConvBlock(\n",
       "      (block): Sequential(\n",
       "        (0): Conv2d(4, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "    )\n",
       "    (1): ConvBlock(\n",
       "      (block): Sequential(\n",
       "        (0): Conv2d(4, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (activation): GReLU(leak=0.1, max=6.0, sub=0.4)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ResBlock(4, norm=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
