{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-30T00:25:20.423936Z",
     "start_time": "2020-12-30T00:25:19.838491Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/hmamin/incendio/incendio/callbacks.py:26: UserWarning: Accio not available.\n",
      "  warnings.warn('Accio not available.')\n"
     ]
    }
   ],
   "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": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T22:57:12.256268Z",
     "start_time": "2020-12-29T22:57:11.894435Z"
    }
   },
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T22:57:12.788559Z",
     "start_time": "2020-12-29T22:57:12.258272Z"
    }
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-01-03T04:18:02.091300Z",
     "start_time": "2021-01-03T04:18:02.048267Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/ipykernel/ipkernel.py:287: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
      "  and should_run_async(code)\n"
     ]
    }
   ],
   "source": [
    "# export\n",
    "from abc import abstractmethod, ABC\n",
    "import copy\n",
    "from einops.layers.torch import Rearrange\n",
    "from functools import partial\n",
    "import numpy as np\n",
    "from operator import add, truediv, sub\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from htools import add_docstring\n",
    "from incendio.core import BaseModel\n",
    "from incendio.data import probabilistic_hash_tensor\n",
    "from incendio.utils import concat, weighted_avg, identity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T22:57:15.652884Z",
     "start_time": "2020-12-29T22:57:15.602809Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/ipykernel/ipkernel.py:287: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
      "  and should_run_async(code)\n",
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/pandas_flavor/register.py:31: UserWarning: registration of accessor <class 'pandas_flavor.register.register_dataframe_method.<locals>.inner.<locals>.AccessorMethod'> under name 'ends' for type <class 'pandas.core.frame.DataFrame'> is overriding a preexisting attribute with the same name.\n",
      "  register_dataframe_accessor(method.__name__)(AccessorMethod)\n",
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/pandas_flavor/register.py:53: UserWarning: registration of accessor <class 'pandas_flavor.register.register_series_method.<locals>.inner.<locals>.AccessorMethod'> under name 'ends' for type <class 'pandas.core.series.Series'> is overriding a preexisting attribute with the same name.\n",
      "  register_series_accessor(method.__name__)(AccessorMethod)\n",
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/pandas_flavor/register.py:31: UserWarning: registration of accessor <class 'pandas_flavor.register.register_dataframe_method.<locals>.inner.<locals>.AccessorMethod'> under name 'filter_by_count' for type <class 'pandas.core.frame.DataFrame'> is overriding a preexisting attribute with the same name.\n",
      "  register_dataframe_accessor(method.__name__)(AccessorMethod)\n",
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/pandas_flavor/register.py:31: UserWarning: registration of accessor <class 'pandas_flavor.register.register_dataframe_method.<locals>.inner.<locals>.AccessorMethod'> under name 'grouped_mode' for type <class 'pandas.core.frame.DataFrame'> is overriding a preexisting attribute with the same name.\n",
      "  register_dataframe_accessor(method.__name__)(AccessorMethod)\n",
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/pandas_flavor/register.py:31: UserWarning: registration of accessor <class 'pandas_flavor.register.register_dataframe_method.<locals>.inner.<locals>.AccessorMethod'> under name 'impute' for type <class 'pandas.core.frame.DataFrame'> is overriding a preexisting attribute with the same name.\n",
      "  register_dataframe_accessor(method.__name__)(AccessorMethod)\n",
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/pandas_flavor/register.py:31: UserWarning: registration of accessor <class 'pandas_flavor.register.register_dataframe_method.<locals>.inner.<locals>.AccessorMethod'> under name 'target_encode' for type <class 'pandas.core.frame.DataFrame'> is overriding a preexisting attribute with the same name.\n",
      "  register_dataframe_accessor(method.__name__)(AccessorMethod)\n",
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/pandas_flavor/register.py:31: UserWarning: registration of accessor <class 'pandas_flavor.register.register_dataframe_method.<locals>.inner.<locals>.AccessorMethod'> under name 'top_categories' for type <class 'pandas.core.frame.DataFrame'> is overriding a preexisting attribute with the same name.\n",
      "  register_dataframe_accessor(method.__name__)(AccessorMethod)\n",
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/pandas_flavor/register.py:53: UserWarning: registration of accessor <class 'pandas_flavor.register.register_series_method.<locals>.inner.<locals>.AccessorMethod'> under name 'vcounts' for type <class 'pandas.core.series.Series'> is overriding a preexisting attribute with the same name.\n",
      "  register_series_accessor(method.__name__)(AccessorMethod)\n",
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/pandas_flavor/register.py:31: UserWarning: registration of accessor <class 'pandas_flavor.register.register_dataframe_method.<locals>.inner.<locals>.AccessorMethod'> under name 'pprint' for type <class 'pandas.core.frame.DataFrame'> is overriding a preexisting attribute with the same name.\n",
      "  register_dataframe_accessor(method.__name__)(AccessorMethod)\n",
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/pandas_flavor/register.py:53: UserWarning: registration of accessor <class 'pandas_flavor.register.register_series_method.<locals>.inner.<locals>.AccessorMethod'> under name 'pprint' for type <class 'pandas.core.series.Series'> is overriding a preexisting attribute with the same name.\n",
      "  register_series_accessor(method.__name__)(AccessorMethod)\n",
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/pandas_flavor/register.py:31: UserWarning: registration of accessor <class 'pandas_flavor.register.register_dataframe_method.<locals>.inner.<locals>.AccessorMethod'> under name 'lambda_sort' for type <class 'pandas.core.frame.DataFrame'> is overriding a preexisting attribute with the same name.\n",
      "  register_dataframe_accessor(method.__name__)(AccessorMethod)\n",
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/pandas_flavor/register.py:53: UserWarning: registration of accessor <class 'pandas_flavor.register.register_series_method.<locals>.inner.<locals>.AccessorMethod'> under name 'lambda_sort' for type <class 'pandas.core.series.Series'> is overriding a preexisting attribute with the same name.\n",
      "  register_series_accessor(method.__name__)(AccessorMethod)\n",
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/pandas_flavor/register.py:31: UserWarning: registration of accessor <class 'pandas_flavor.register.register_dataframe_method.<locals>.inner.<locals>.AccessorMethod'> under name 'coalesce' for type <class 'pandas.core.frame.DataFrame'> is overriding a preexisting attribute with the same name.\n",
      "  register_dataframe_accessor(method.__name__)(AccessorMethod)\n",
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/pandas_flavor/register.py:53: UserWarning: registration of accessor <class 'pandas_flavor.register.register_series_method.<locals>.inner.<locals>.AccessorMethod'> under name 'stringify' for type <class 'pandas.core.series.Series'> is overriding a preexisting attribute with the same name.\n",
      "  register_series_accessor(method.__name__)(AccessorMethod)\n",
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/pandas_flavor/register.py:53: UserWarning: registration of accessor <class 'pandas_flavor.register.register_series_method.<locals>.inner.<locals>.AccessorMethod'> under name 'is_list_col' for type <class 'pandas.core.series.Series'> is overriding a preexisting attribute with the same name.\n",
      "  register_series_accessor(method.__name__)(AccessorMethod)\n"
     ]
    }
   ],
   "source": [
    "# Used for testing only.\n",
    "from collections import defaultdict, Counter\n",
    "from itertools import chain\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "from htools import assert_raises, InvalidArgumentError, smap\n",
    "from incendio.data import probabilistic_hash_item\n",
    "import pandas_htools"
   ]
  },
  {
   "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. This is an object, not a class.\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",
    "        return 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": [
    "conv = ConvBlock(3, 5, norm=False)\n",
    "conv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 2, 2])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.rand(2, 3, 4, 4)\n",
    "conv(x).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class ResBlock(nn.Module):\n",
    "\n",
    "    def __init__(self, c_in, kernel_size=3, norm=True, activation=JRelu, \n",
    "                 stride=1, padding=1, skip_size=2, **kwargs):\n",
    "        \"\"\"Residual block using 2D convolutional layers. Note that kernel_size,\n",
    "        stride, and pad must be selected such that the height and width of \n",
    "        the input remain the same.\n",
    "\n",
    "        Parameters\n",
    "        -----------\n",
    "        c_in: int\n",
    "            # of input channels.\n",
    "        kernel_size: int\n",
    "            Size of filter used in convolution. Default 3 (which becomes 3x3).\n",
    "        norm: bool\n",
    "            Specifies whether to include a batch norm layer after each conv\n",
    "            layer.\n",
    "        activation: callable\n",
    "            Activation function to use.\n",
    "        stride: int\n",
    "            # of pixels the filter moves between each convolution. Default 1.\n",
    "        padding: 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",
    "        kwargs: any\n",
    "            Additional kwargs to pass to ConvBlock which will in turn pass them\n",
    "            to Conv2d. If you accidentally pass in a 'c_out', it will be \n",
    "            removed since we need all dimensions to remain unchanged.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        # Ensure we don't accidentally pass in a different c_out.\n",
    "        kwargs.pop('c_out', None)\n",
    "        self.skip_size = skip_size\n",
    "        self.layers = nn.ModuleList([\n",
    "            ConvBlock(c_in, c_in, kernel_size=kernel_size, norm=norm, \n",
    "                      activation=None, stride=stride, padding=padding, \n",
    "                      **kwargs)\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "@add_docstring(nn.Conv2d)\n",
    "class ReflectionPaddedConv2d(nn.Module):\n",
    "    \"\"\"Conv2d only allows padding_mode of `zeros` or `circular`. This\n",
    "    layer is a quick way for us to use reflection padding.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, in_channels, out_channels, padding=1, \n",
    "                 kernel_size=3, **kwargs):\n",
    "        \"\"\"Do not specify a padding mode.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        if 'padding_mode' in kwargs: \n",
    "            raise InvalidArgumentError('Remove `padding_mode` from arguments.')\n",
    "        self.reflect = nn.ReflectionPad2d(padding)\n",
    "        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size,\n",
    "                              padding=0)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.reflect(x)\n",
    "        return self.conv(x)"
   ]
  },
  {
   "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": {
      "text/plain": [
       "ReflectionPaddedConv2d(\n",
       "  (reflect): ReflectionPad2d((2, 2, 2, 2))\n",
       "  (conv): Conv2d(3, 3, kernel_size=(1, 1), stride=(1, 1))\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rconv = ReflectionPaddedConv2d(3, 3, kernel_size=1, padding=2)\n",
    "rconv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 5016 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "x = torch.randint(255, (1, 3, 3, 3)).float()\n",
    "show_img(x[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 4168 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "x2 = rconv.reflect(x)\n",
    "show_img(x2[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "As expected, got InvalidArgumentError(Remove `padding_mode` from arguments.).\n"
     ]
    }
   ],
   "source": [
    "# Tests\n",
    "assert nn.Conv2d.__doc__ in ReflectionPaddedConv2d.__doc__\n",
    "\n",
    "with assert_raises(InvalidArgumentError):\n",
    "    ReflectionPaddedConv2d(3, 3, padding_mode='zeros')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T22:57:23.822261Z",
     "start_time": "2020-12-29T22:57:23.780361Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class SmoothSoftmaxBase(nn.Module):\n",
    "    \"\"\"Parent class of SmoothSoftmax and SmoothLogSoftmax (softmax or log\n",
    "    softmax with temperature baked in). There shouldn't be a need to\n",
    "    instantiate this class directly.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, log=False, temperature='auto', dim=-1):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        log: bool\n",
    "            If True, use log softmax (if this is the last activation in a\n",
    "            network, it can be followed by nn.NLLLoss). If False, use softmax\n",
    "            (this is more useful if you're doing something attention-related:\n",
    "            no standard torch loss functions expect softmax outputs). This\n",
    "            argument is usually passed implicitly by the higher level interface\n",
    "            provided by the child classes.\n",
    "        temperature: float or str\n",
    "            If a float, this is the temperature to divide activations by before\n",
    "            applying the softmax. Values larger than 1 soften the distribution\n",
    "            while values between 0 and 1 sharpen it. If str ('auto'), this will\n",
    "            compute the square root of the last dimension of x's shape the\n",
    "            first time the forward method is called and use that for subsequent\n",
    "            calls.\n",
    "        dim: int\n",
    "            The dimension to compute the softmax over.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.temperature = None if temperature == 'auto' else temperature\n",
    "        self.act = nn.LogSoftmax(dim=dim) if log else nn.Softmax(dim=dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        x: torch.float\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        torch.float: Same shape as x.\n",
    "        \"\"\"\n",
    "        # Slightly odd way to do this but we're trying to avoid an extra if\n",
    "        # statement because temperature only needs to be set once and we could\n",
    "        # plausibly call this method millions of times during training.\n",
    "        try:\n",
    "            return self.act(x.div(self.temperature))\n",
    "        except TypeError:\n",
    "            self.temperature = np.sqrt(x.shape[-1])\n",
    "            return self.forward(x)\n",
    "        except Exception as e:\n",
    "            raise e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T22:57:24.880587Z",
     "start_time": "2020-12-29T22:57:24.837070Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class SmoothSoftmax(SmoothSoftmaxBase):\n",
    "\n",
    "    def __init__(self, temperature='auto', dim=-1):\n",
    "        super().__init__(log=False, temperature=temperature, dim=dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T22:57:25.245163Z",
     "start_time": "2020-12-29T22:57:25.189414Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class SmoothLogSoftmax(SmoothSoftmaxBase):\n",
    "\n",
    "    def __init__(self, temperature='auto', dim=-1):\n",
    "        super().__init__(log=True, temperature=temperature, dim=dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T22:58:15.139326Z",
     "start_time": "2020-12-29T22:58:15.097502Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/ipykernel/ipkernel.py:287: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
      "  and should_run_async(code)\n"
     ]
    }
   ],
   "source": [
    "# export\n",
    "class SpatialSoftmax(nn.Module):\n",
    "    \"\"\"Apply softmax over the height and width dimensions of a batch of image\n",
    "    tensors (or image-like tensors). Concretely, inputs will usually have\n",
    "    shape (batch size, channels, height, width), while outputs will have the\n",
    "    same shape but values for each feature map will now sum to 1. Essentially,\n",
    "    we now have a heatmap of what region in each image to focus on.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, temperature='auto', log=False):\n",
    "        super().__init__()\n",
    "        cls = SmoothLogSoftmax if log else SmoothSoftmax\n",
    "        self.act = cls(temperature)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # Should work on any tensor with shape (bs, ..., h, w).\n",
    "        flattened = self.act(x.view(*x.shape[:-2], -1))\n",
    "        return flattened.view(*x.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class Dropin(nn.Module):\n",
    "    \"\"\"Additive dropout. This injects small amounts of noise into a model\n",
    "    in the form of randomly generated floats from a zero-centered\n",
    "    gaussian distribution (variance can be adjusted). This does nothing \n",
    "    in eval mode. Unlike Dropout, this does not scale weights during \n",
    "    training since it does not bias them in any direction.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, scale=.5):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        scale: float\n",
    "            Used to scale the magnitude of the random noise. Keep in mind \n",
    "            that the scalar term is square rooted, so the relationship\n",
    "            will not be linear. Relatively large values (e.g. 1.0) will have\n",
    "            a stronger regularizing effect, while small values (e.g. 0.1)\n",
    "            will have a slight regularizing effect. There is no max value\n",
    "            enforced, so it's up to the user to select a reasonable value.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.scale = scale\n",
    "        \n",
    "    def forward(self, x):\n",
    "        if not self.training:\n",
    "            return x\n",
    "        \n",
    "        # Storing noise allows us to run diagnostics.\n",
    "        self.noise = torch.randn_like(x) * np.sqrt(self.scale / x.shape[-1])\n",
    "        return x + self.noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.drop = Dropin()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.drop(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Net()\n",
    "x = torch.randn(8, 128, 128, 3)\n",
    "assert np.corrcoef(net(x).flatten(), x.flatten())[0][1] > .9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net.eval()\n",
    "assert torch.eq(net(x), x).all()\n",
    "assert not net.drop.training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def simulate_activation_stats(scale=1.0, trials=10_000):\n",
    "    act_stats = defaultdict(list)\n",
    "    noise_stats = defaultdict(list)\n",
    "    \n",
    "    drop = Dropin(scale)\n",
    "    for _ in range(trials):\n",
    "        x = torch.randn(3, 4, dtype=torch.float)\n",
    "        z = drop(x)\n",
    "        noise = drop.noise\n",
    "        noise_stats['mean'].append(noise.mean())\n",
    "        noise_stats['std'].append(noise.std())\n",
    "        noise_stats['act_corr'].append(\n",
    "            np.corrcoef(z.flatten(), noise.flatten())[0][1]\n",
    "        )\n",
    "        \n",
    "        act_stats['mean'].append(z.mean())\n",
    "        act_stats['std'].append(z.std())\n",
    "        act_stats['x_corr'].append(\n",
    "            np.corrcoef(z.flatten(), x.flatten())[0][1]\n",
    "        )\n",
    "\n",
    "    return pd.DataFrame(dict(\n",
    "        act={k: np.mean(v).round(4) for k, v in act_stats.items()}, \n",
    "        noise={k: np.mean(v).round(4) for k, v in noise_stats.items()}\n",
    "    ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " 10\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>act</th>\n",
       "      <th>noise</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.0132</td>\n",
       "      <td>0.0094</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>1.8189</td>\n",
       "      <td>1.5192</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x_corr</th>\n",
       "      <td>0.5324</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>act_corr</th>\n",
       "      <td>NaN</td>\n",
       "      <td>0.8304</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " 1\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>act</th>\n",
       "      <th>noise</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>-0.0141</td>\n",
       "      <td>0.0034</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>1.0921</td>\n",
       "      <td>0.4870</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x_corr</th>\n",
       "      <td>0.8855</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>act_corr</th>\n",
       "      <td>NaN</td>\n",
       "      <td>0.4282</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " 0.75\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>act</th>\n",
       "      <th>noise</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>-0.0015</td>\n",
       "      <td>0.0022</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>1.0633</td>\n",
       "      <td>0.4240</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x_corr</th>\n",
       "      <td>0.9100</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>act_corr</th>\n",
       "      <td>NaN</td>\n",
       "      <td>0.3899</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " 0.5\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>act</th>\n",
       "      <th>noise</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.0107</td>\n",
       "      <td>0.0008</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>1.0558</td>\n",
       "      <td>0.3442</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x_corr</th>\n",
       "      <td>0.9409</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>act_corr</th>\n",
       "      <td>NaN</td>\n",
       "      <td>0.3235</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " 0.25\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>act</th>\n",
       "      <th>noise</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.0098</td>\n",
       "      <td>-0.0057</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>1.0013</td>\n",
       "      <td>0.2461</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x_corr</th>\n",
       "      <td>0.9667</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>act_corr</th>\n",
       "      <td>NaN</td>\n",
       "      <td>0.2298</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " 0.1\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>act</th>\n",
       "      <th>noise</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>-0.0057</td>\n",
       "      <td>-0.0014</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.9969</td>\n",
       "      <td>0.1533</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x_corr</th>\n",
       "      <td>0.9868</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>act_corr</th>\n",
       "      <td>NaN</td>\n",
       "      <td>0.1394</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for scale in [10, 1, .75, .5, .25, .1]:\n",
    "    print('\\n', scale)\n",
    "    simulate_activation_stats(scale, 1_000).pprint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class LinearSkipBlock(nn.Module):\n",
    "    \"\"\"This lets us easily create residual block equivalents with linear\n",
    "    layers.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, x_dim, layer_dims, op, activation=mish):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        x_dim: int\n",
    "            Size of input tensor.\n",
    "        layer_dims: Iterable[int]\n",
    "            Size of each layer. The length of this list will be the skip size\n",
    "            (2 is probably a reasonable starting point).\n",
    "        op: function\n",
    "            This will be called on the input x and the processed x in the\n",
    "            forward method. This is a concatenation for dense blocks and an\n",
    "            addition for residual blocks, but any operation is possible.\n",
    "        activation: callable\n",
    "            Activation function or callable class. This will be applied after\n",
    "            each layer. The final activation is applied after the `op` \n",
    "            function.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.skip_size = len(layer_dims)\n",
    "        self.activation = activation\n",
    "        self.layers = nn.ModuleList([nn.Linear(d_in, d_out) for d_in, d_out\n",
    "                                     in zip([x_dim]+list(layer_dims), \n",
    "                                            layer_dims)])\n",
    "        self.op = op\n",
    "            \n",
    "    def forward(self, x):\n",
    "        out = x\n",
    "        for i, layer in enumerate(self.layers, 1):\n",
    "            out = layer(out)\n",
    "            if i < self.skip_size: out = self.activation(out)\n",
    "        return self.activation(self.op(x, out))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class LinearResBlock(LinearSkipBlock):\n",
    "    \"\"\"Equivalent of ResNet block with linear layers.\"\"\"\n",
    "    \n",
    "    def __init__(self, x_dim, hidden_dims, activation=mish):\n",
    "        if hidden_dims[-1] != x_dim:\n",
    "            raise InvalidArgumentError(\n",
    "                'Last hidden dimension must match input dimension.'\n",
    "            )\n",
    "        super().__init__(x_dim, hidden_dims, add, activation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class LinearDenseBlock(LinearSkipBlock):\n",
    "    \"\"\"Equivalent of DenseNet block with linear layers.\"\"\"\n",
    "    \n",
    "    def __init__(self, x_dim, hidden_dims, activation=mish):\n",
    "        super().__init__(x_dim, hidden_dims, concat, activation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class WeightedLinearResBlock(LinearSkipBlock):\n",
    "    \"\"\"Like a LinearResBlock but takes a weighted average of the input and \n",
    "    output rather than adding them. Addition gives them equal weight and we \n",
    "    may want to weight the output more heavily.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, x_dim, hidden_dims, weights=(.25, .75),\n",
    "                 activation=mish):\n",
    "        super().__init__(x_dim, hidden_dims, \n",
    "                         partial(weighted_avg, weights=list(weights)),\n",
    "                         activation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class SkipConnection(nn.Module):\n",
    "    \"\"\"More generalized version of skip connection. Eventually maybe rewrite\n",
    "    various res/dense/weighted conv blocks with this.\n",
    "    \n",
    "    Examples\n",
    "    --------\n",
    "    >> x = torch.randn(3, 4)\n",
    "    >> dense = nn.Linear(4, 2)\n",
    "    >> dense(x).shape\n",
    "    \n",
    "    torch.Size([3, 2])\n",
    "    \n",
    "    >> skip = SkipConnection(dense, op='cat')\n",
    "    >> skip(x).shape\n",
    "    \n",
    "    torch.Size([3, 6])\n",
    "    \n",
    "    >> skip = SkipConnection(dense, op='add')\n",
    "    >> skip(x).shape\n",
    "    \n",
    "    RuntimeError: The size of tensor a (4) must match the size of tensor b (2)\n",
    "    at non-singleton dimension 1\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, block, op='add', input_weight=None):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        block: nn.Module\n",
    "            A torch layer/model that takes in some input x (and optionally\n",
    "            other args/kwargs) and performs some computations on it. When \n",
    "            using op='add', this should output a tensor with the same shape\n",
    "            as its first input.\n",
    "        op: str\n",
    "            One of ('add', 'cat', 'weighted_avg'). This determines how the\n",
    "            input will be attached to the output. If you choose 'cat',\n",
    "            concatenation will occur over the last axis and input will precede\n",
    "            output.\n",
    "        input_weight: float or None\n",
    "            If op='weighted_avg', you must provide a float in (0, 1) that\n",
    "            determines how heavily to weight the input x. For example, 0.2\n",
    "            means the output of `block` will be much more heavily weighted\n",
    "            than the input tensor, while 0.5 is equivalent to computing the \n",
    "            mean (and in most cases is essentially equivalent to computing\n",
    "            the sum).\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.block = block\n",
    "        if op == 'add':\n",
    "            self.op = torch.add\n",
    "        elif op == 'cat':\n",
    "            self.op = self._cat\n",
    "        elif op == 'weighted_avg':\n",
    "            if input_weight is None or input_weight <= 0 or input_weight >= 1:\n",
    "                raise ValueError('input_weight must be a float in (0, 1) '\n",
    "                                 'when op=\"weighted\".')\n",
    "            self.weights = input_weight, 1-input_weight\n",
    "            self.op = self._weighted_avg\n",
    "        else:\n",
    "            raise ValueError('op must be in (\"add\", \"cat\", \"weighted_avg\").')\n",
    "        \n",
    "    def forward(self, x, *args, **kwargs):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        x: torch.Tensor\n",
    "            This first item is considered to be the input which will be \n",
    "            combined with the output of self.block.\n",
    "        args, kwargs: any\n",
    "            Additional args will be forwarded to self.block.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        torch.Tensor: Should have same shape as x unless you're making use of\n",
    "        broadcasting, which should rarely be needed here.\n",
    "        \"\"\"\n",
    "        return self.op(x, self.block(x, *args, **kwargs))\n",
    "        \n",
    "    @staticmethod\n",
    "    def _cat(x1, x2):\n",
    "        \"\"\"Wrapper since torch.cat has a different interface than torch.add\n",
    "        (list of args vs. *args).\n",
    "        \"\"\"\n",
    "        return torch.cat([x1, x2], dim=-1)\n",
    "    \n",
    "    def _weighted_avg(x1, x2):\n",
    "        \"\"\"In our use case, the first tensor will be the original input tensor\n",
    "        and the second will be the output of self.block.\n",
    "        \"\"\"\n",
    "        return self.weights[0]*x1 + self.weights[1]*x2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Embeddings and Encodings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-21T21:18:01.876862Z",
     "start_time": "2020-11-21T21:18:01.837126Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def trunc_normal_(x, mean=0.0, std=1.0):\n",
    "    \"\"\"Ported from fastai to remove dependency: \n",
    "    \n",
    "    Truncated normal initialization.\n",
    "    From https://discuss.pytorch.org/t/implementing-truncated-normal-initializer/4778/12\n",
    "    \"\"\"\n",
    "    return x.normal_().fmod_(2).mul_(std).add_(mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-21T21:31:57.191450Z",
     "start_time": "2020-11-21T21:31:57.143277Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class InitializedEmbedding(nn.Embedding):\n",
    "    \"\"\"Same as nn.Embedding but with truncated normal initialization. This\n",
    "    also differs from fastai's Embedding class in that it allows padding.\n",
    "    \"\"\"\n",
    "        \n",
    "    def reset_parameters(self):\n",
    "        with torch.no_grad():\n",
    "            trunc_normal_(self.weight, std=.01)\n",
    "            if self.padding_idx is not None:\n",
    "                torch.zero_(self.weight[self.padding_idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-21T21:31:58.004885Z",
     "start_time": "2020-11-21T21:31:57.962186Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[ 0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0031, -0.0033, -0.0013],\n",
       "        [ 0.0076, -0.0107,  0.0008],\n",
       "        [ 0.0119, -0.0014,  0.0027]], requires_grad=True)"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "InitializedEmbedding(4, 3, 0).weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[ 0.0038, -0.0103, -0.0074],\n",
       "        [ 0.0104, -0.0036,  0.0039],\n",
       "        [-0.0007,  0.0008,  0.0191],\n",
       "        [ 0.0000,  0.0000,  0.0000]], requires_grad=True)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "InitializedEmbedding(4, 3, 3).weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[-0.0082,  0.0058, -0.0055],\n",
       "        [-0.0104, -0.0144,  0.0036],\n",
       "        [ 0.0028, -0.0020, -0.0009],\n",
       "        [ 0.0046, -0.0088, -0.0026]], requires_grad=True)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "InitializedEmbedding(4, 3).weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-21T21:34:15.848366Z",
     "start_time": "2020-11-21T21:34:15.789283Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class BloomEmbedding(nn.Module):\n",
    "    \"\"\"Bloom Embedding layer for memory-efficient word representations.\n",
    "    Each word is encoded by a combination of rows of the embedding\n",
    "    matrix. The number of rows can therefore be far lower than the number\n",
    "    of words in our vocabulary while still providing unique representations.\n",
    "    The reduction in rows allows us to use memory in other ways: a larger\n",
    "    embedding dimension, more or larger layers after the embedding,\n",
    "    larger batch sizes, etc.\n",
    "    \n",
    "    Note that if hashing is done in the Dataset, we could use a simple\n",
    "    nn.EmbeddingBag to achieve the same thing. Many users have reported \n",
    "    poor performance with this layer though (especially on CPU, but in some\n",
    "    cases on GPU) so I stick with the standard Embedding. We also bake in\n",
    "    the truncated normal intialization provided by fastai, with a slight tweak\n",
    "    to allow a row for padding.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, n_emb=251, emb_dim=100, n_hashes=4, padding_idx=0,\n",
    "                 pre_hashed=False):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        n_emb: int\n",
    "            Number of rows to create in the embedding matrix. A prime\n",
    "            number is recommended. Lower numbers will be more \n",
    "            memory-efficient but increase the chances of collisions.\n",
    "        emb_dim: int\n",
    "            Size of each embedding. If emb_dim=100, each word will\n",
    "            be represented by a 100-dimensional vector.\n",
    "        n_hashes: int\n",
    "            This determines the number of hashes that will be taken\n",
    "            for each word index, and as a result, the number of rows\n",
    "            that will be summed to create each unique representation.\n",
    "            The higher the number, the lower the chances of a collision.\n",
    "        padding_idx: int or None\n",
    "            If an integer is provided, this will set aside the corresponding\n",
    "            row in the embedding matrix as a vector of zeros. If None, no\n",
    "            padding vector will be allocated.\n",
    "        pre_hashed: bool\n",
    "            Pass in True if the input tensor will already be hashed by the\n",
    "            time it enters this layer (you may prefer pre-compute the hashes \n",
    "            in the Dataset to save computation time during training). In this\n",
    "            scenario, the layer is a simple embedding bag with mode \"sum\". \n",
    "            Pass in False if the inputs will be word indices that have not yet\n",
    "            been hashed. In this case, hashing will be done inside the \n",
    "            `forward` call.\n",
    "            \n",
    "        Suggested values for a vocab size of ~30,000:\n",
    "        \n",
    "        | n_emb | n_hashes | unique combos |\n",
    "        |-------|----------|---------------|\n",
    "        | 127   | 5        | 29,998        |\n",
    "        | 251   | 4        | 29,996        |\n",
    "        | 997   | 3        | 29,997        |\n",
    "        | 5,003 | 2        | 29,969        |\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.n_emb = n_emb\n",
    "        self.emb = InitializedEmbedding(n_emb, emb_dim, padding_idx)\n",
    "        self.n_hashes = n_hashes\n",
    "        self.pad_idx = padding_idx\n",
    "        self.pre_hashed = pre_hashed\n",
    "        self.process_fn = identity if pre_hashed else \\\n",
    "            partial(probabilistic_hash_tensor, n_buckets=n_emb,\n",
    "                    n_hashes=n_hashes, pad_idx=padding_idx)\n",
    "        # Makes interface consistent with nn.Embedding. Don't change name.\n",
    "        self.embedding_dim = self.emb.embedding_dim\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        x: torch.LongTensor\n",
    "            Input tensor of word indices (bs x seq_len) if pre_hashed is \n",
    "            False. Hashed indices (bs x seq_len x n_hashes) if pre_hashed is \n",
    "            False.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        torch.FloatTensor: Words encoded with combination of embeddings.\n",
    "            (bs x seq_len x emb_dim)\n",
    "        \"\"\"\n",
    "        # If not pre-hashed: (bs, seq_len) -> hash -> (bs, seq_len, n_hashes)\n",
    "        hashed = self.process_fn(x)\n",
    "        # (bs, seq_len, n_hashes, emb_dim) -> sum -> (bs, seq_len, emb_dim)\n",
    "        return self.emb(hashed).sum(-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-21T21:31:16.317372Z",
     "start_time": "2020-11-21T21:31:16.276390Z"
    }
   },
   "outputs": [],
   "source": [
    "class Data(Dataset):\n",
    "    \n",
    "    def __init__(self, sentences, labels, seq_len):\n",
    "        x = [s.split(' ') for s in sentences]\n",
    "        self.w2i = self.make_w2i(x)\n",
    "        self.seq_len = seq_len\n",
    "        self.x = self.encode(x)\n",
    "        self.y = torch.tensor(labels)\n",
    "        \n",
    "    def __getitem__(self, i):\n",
    "        return self.x[i], self.y[i]\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.y)\n",
    "    \n",
    "    def make_w2i(self, tok_rows):\n",
    "        return {k: i for i, (k, v) in \n",
    "                enumerate(Counter(chain(*tok_rows)).most_common(), 1)}\n",
    "    \n",
    "    def encode(self, tok_rows):\n",
    "        enc = np.zeros((len(tok_rows), self.seq_len), dtype=int)\n",
    "        for i, row in enumerate(tok_rows):\n",
    "            trunc = [self.w2i.get(w, 0) for w in row[:self.seq_len]]\n",
    "            enc[i, :len(trunc)] = trunc\n",
    "        return torch.tensor(enc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-21T21:31:16.998818Z",
     "start_time": "2020-11-21T21:31:16.957720Z"
    }
   },
   "outputs": [],
   "source": [
    "sents = [\n",
    "    'I walked to the store so I hope it is not closed.',\n",
    "    'The theater is closed today and the sky is grey.',\n",
    "    'His dog is brown while hers is grey.'\n",
    "]\n",
    "labels = [0, 1, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-21T21:31:17.702944Z",
     "start_time": "2020-11-21T21:31:17.645462Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([13, 14,  1, 15, 16, 17,  3, 18,  1,  4]), tensor(1))"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds = Data(sents, labels, 10)\n",
    "ds[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-21T21:31:18.647254Z",
     "start_time": "2020-11-21T21:31:18.589002Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 2,  5,  6,  3,  7,  8,  2,  9, 10,  1],\n",
       "         [13, 14,  1, 15, 16, 17,  3, 18,  1,  4],\n",
       "         [19, 20,  1, 21, 22, 23,  1,  4,  0,  0]]),\n",
       " tensor([0, 1, 1]))"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dl = DataLoader(ds, batch_size=3)\n",
    "x, y = next(iter(dl))\n",
    "x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-21T21:31:19.235283Z",
     "start_time": "2020-11-21T21:31:19.188061Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 2,  5,  6,  3,  7,  8,  2,  9, 10,  1],\n",
       "         [13, 14,  1, 15, 16, 17,  3, 18,  1,  4],\n",
       "         [19, 20,  1, 21, 22, 23,  1,  4,  0,  0]]),\n",
       " tensor([0, 1, 1]))"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x, y = next(iter(dl))\n",
    "x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-21T21:31:20.247703Z",
     "start_time": "2020-11-21T21:31:20.197976Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
       "        [-1.3866e-02,  2.7083e-03,  1.9108e-03,  2.3947e-03],\n",
       "        [ 5.0557e-03, -4.8774e-03,  1.5206e-03,  1.0080e-03],\n",
       "        [ 4.5406e-05,  2.1932e-03, -9.5886e-03,  4.7936e-03],\n",
       "        [-9.7800e-03,  1.9145e-02,  2.4422e-03,  1.2713e-02],\n",
       "        [ 1.2850e-02,  4.0345e-03, -1.9255e-02, -2.5600e-03],\n",
       "        [ 5.7765e-03,  1.4253e-02,  1.8160e-02, -1.6686e-02],\n",
       "        [ 1.2343e-02,  2.8021e-03,  9.5432e-04,  4.2866e-03],\n",
       "        [ 6.4279e-03, -3.4935e-03, -1.9902e-03, -9.9574e-03],\n",
       "        [ 9.7122e-04, -6.8190e-03, -1.2612e-02, -1.6921e-03],\n",
       "        [-1.1588e-02, -4.2316e-03, -9.5648e-03, -6.5988e-03]],\n",
       "       requires_grad=True)"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "be = BloomEmbedding(11, 4)\n",
    "be.emb.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 2,  5,  6,  3,  7,  8,  2,  9, 10,  1],\n",
       "        [13, 14,  1, 15, 16, 17,  3, 18,  1,  4],\n",
       "        [19, 20,  1, 21, 22, 23,  1,  4,  0,  0]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 10, 4])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# (bs x seq_len) -> (bs -> seq_len -> emb_size)\n",
    "y = be(x)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0363, -0.0098,  0.0052,  0.0023],\n",
       "        [ 0.0337, -0.0130, -0.0147, -0.0183],\n",
       "        [ 0.0075, -0.0316,  0.0120,  0.0236],\n",
       "        [ 0.0304, -0.0328,  0.0388, -0.0240],\n",
       "        [ 0.0144,  0.0004, -0.0022,  0.0019],\n",
       "        [ 0.0084,  0.0117,  0.0089, -0.0284],\n",
       "        [ 0.0363, -0.0098,  0.0052,  0.0023],\n",
       "        [ 0.0177, -0.0087,  0.0344, -0.0139],\n",
       "        [ 0.0272, -0.0108, -0.0036,  0.0130],\n",
       "        [ 0.0035, -0.0162,  0.0048,  0.0260]], grad_fn=<SelectBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below, we show by step how to get from x to y. This is meant to demonstrate the basic mechanism, not to show how PyTorch actually implements this under the hood. Let's look at a single row of x, corresponding to 1 sentence where each word is mapped to its index in the vocabulary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 2,  5,  6,  3,  7,  8,  2,  9, 10,  1])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we hash each item."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[8, 2, 7, 8],\n",
       " [2, 8, 1, 2],\n",
       " [6, 6, 10, 10],\n",
       " [10, 5, 5, 5],\n",
       " [6, 9, 7, 2],\n",
       " [5, 9, 4, 0],\n",
       " [8, 2, 7, 8],\n",
       " [5, 10, 8, 9],\n",
       " [7, 8, 6, 2],\n",
       " [6, 10, 6, 0]]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hashed = [probabilistic_hash_item(i.item(), 11, int, 4) for i in x[0]]\n",
    "hashed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then use each row of hashed integers to index into the embedding weight matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10, 4, 4])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.0089,  0.0007,  0.0076,  0.0034],\n",
       "         [ 0.0097,  0.0003, -0.0098, -0.0082],\n",
       "         [ 0.0089, -0.0114, -0.0001,  0.0037],\n",
       "         [ 0.0089,  0.0007,  0.0076,  0.0034]],\n",
       "\n",
       "        [[ 0.0097,  0.0003, -0.0098, -0.0082],\n",
       "         [ 0.0089,  0.0007,  0.0076,  0.0034],\n",
       "         [ 0.0054, -0.0142, -0.0027, -0.0052],\n",
       "         [ 0.0097,  0.0003, -0.0098, -0.0082]]], grad_fn=<SliceBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output = []\n",
    "for row in hashed:\n",
    "    row_out = be.emb.weight[row]\n",
    "    output.append(row_out)\n",
    "output = torch.stack(output)\n",
    "print(output.shape)\n",
    "output[:2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we sum up the embedding rows. Above, each word is represented by four rows of the embedding matrix. After summing, we get a single vector for each word."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0363, -0.0098,  0.0052,  0.0023],\n",
       "        [ 0.0337, -0.0130, -0.0147, -0.0183],\n",
       "        [ 0.0075, -0.0316,  0.0120,  0.0236],\n",
       "        [ 0.0304, -0.0328,  0.0388, -0.0240],\n",
       "        [ 0.0144,  0.0004, -0.0022,  0.0019],\n",
       "        [ 0.0084,  0.0117,  0.0089, -0.0284],\n",
       "        [ 0.0363, -0.0098,  0.0052,  0.0023],\n",
       "        [ 0.0177, -0.0087,  0.0344, -0.0139],\n",
       "        [ 0.0272, -0.0108, -0.0036,  0.0130],\n",
       "        [ 0.0035, -0.0162,  0.0048,  0.0260]], grad_fn=<SumBackward1>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output = output.sum(-2)\n",
    "output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " Notice that the values now match the output of our embedding layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert torch.isclose(output, y[0]).all()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Axial encodings are intended to work as positional embeddings for transformer-like architectures. It's possible they could work for word embeddings as well, similar to our use of Bloom embeddings. However, the standard version of axial encodings results in similar vectors for adjacent indices - this makes some sense for positional indices, but for word indices it might require some additional preprocessing. For example, we could compress word embeddings down to 1 dimension and sort them, or simply sort by number of occurrences in our corpus which could be considered to be doing the same thing. Large chunks of the outputs vectors will be shared among different inputs, whereas Bloom embeddings seem like they would have a greater capacity to avoid this issue."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-21T21:35:16.765623Z",
     "start_time": "2020-11-21T21:35:16.717534Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class AxialEncoding(nn.Module):\n",
    "    \"\"\"Axial encodings. These are intended to encode position in a sequence\n",
    "    (e.g. index in a sentence). It's possible we could adapt these for use as\n",
    "    word embeddings but this would likely require some experimentation (for\n",
    "    example, words would likely need to be sorted in a thoughtful manner \n",
    "    (e.g. pre-trained embeddings compressed to 1D?) since adjacent inputs will\n",
    "    share half of their encodings).\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, vocab_dim, emb_dim, pad_idx=None):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        vocab_dim: int\n",
    "            Number of words in vocab (or max sequence length if being used for\n",
    "            positional encodings).\n",
    "        emb_dim: int\n",
    "            Size of embedding vectors (often numbers like 50, 100, 300).\n",
    "        pad_idx: int or None\n",
    "            If necessary, pass in an integer to represent padding. Otherwise\n",
    "            no rows are reserved for padding.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        if emb_dim % 2 != 0:\n",
    "            raise ValueError('emb_dim must be an even number.')\n",
    "            \n",
    "        self.v = self._decompose_mult(vocab_dim)\n",
    "        self.e = self._decompose_add(emb_dim)\n",
    "        self.emb = nn.ModuleList(InitializedEmbedding(self.v, self.e, pad_idx) \n",
    "                                 for _ in range(2))\n",
    "        # Makes interface consistent with nn.Embedding. Don't change name.\n",
    "        self.embedding_dim = self.e * 2\n",
    "    \n",
    "    def _decompose_mult(self, dim):\n",
    "        return int(np.ceil(np.sqrt(dim)))\n",
    "    \n",
    "    def _decompose_add(self, dim):\n",
    "        return int(np.ceil(dim / 2))\n",
    "    \n",
    "    def forward(self, idx):\n",
    "        return torch.cat([self.emb[0](idx%self.v), self.emb[1](idx//self.v)], \n",
    "                         dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-21T21:35:24.049373Z",
     "start_time": "2020-11-21T21:35:23.953139Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class MultiAxialEncoding(nn.Module):\n",
    "    \"\"\"Adapted axial encodings to allow for more than 2 embedding matrices.\n",
    "    These are intended to encode position in a sequence (e.g. index in a \n",
    "    sentence) but might work as word embeddings. This version may be better \n",
    "    suited for that use case because using more blocks results in fewer shared\n",
    "    numbers in the output vectors of adjacent inputs.\n",
    "    \n",
    "    Some experimentation is still required for this use case (for\n",
    "    example, words would likely need to be sorted in a thoughtful manner \n",
    "    (e.g. pre-trained embeddings compressed to 1D?) since adjacent inputs will\n",
    "    share half of their encodings).\n",
    "    \n",
    "    I made this separate from AxialEncoding (at least for now) since I made a\n",
    "    few tweaks to the original design to make this possible and I wanted to \n",
    "    preserve the option to use the simpler, well-tested method \n",
    "    (AxialEncoding). Here, we use a probabilistic hashing scheme to map each\n",
    "    input to multiple embedding rows, while the original design uses \n",
    "    x%v and x//v.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, vocab_dim, emb_dim, n_blocks=2, pre_hashed=False, \n",
    "                 pad_idx=None):\n",
    "        super().__init__()\n",
    "        # Must set n_blocks before computing v or e.\n",
    "        self.n_blocks = n_blocks\n",
    "        self.v = self._decompose_mult(vocab_dim)\n",
    "        self.e = self._decompose_add(emb_dim)\n",
    "        self.pre_hashed = pre_hashed\n",
    "        # Must set emb blocks before defining process_fn.\n",
    "        self.emb = nn.ModuleList(InitializedEmbedding(self.v, self.e, pad_idx) \n",
    "                                 for _ in range(n_blocks))\n",
    "        self.process_fn = identity if pre_hashed else \\\n",
    "            partial(probabilistic_hash_tensor, n_buckets=self.v, \n",
    "                    n_hashes=len(self.emb), pad_idx=pad_idx)\n",
    "        # Makes interface consistent with nn.Embedding. Don't change name.\n",
    "        self.embedding_dim = self.e * self.n_blocks\n",
    "    \n",
    "    def _decompose_mult(self, dim):\n",
    "        return int(np.ceil(dim ** (1 / self.n_blocks)))\n",
    "    \n",
    "    def _decompose_add(self, dim):\n",
    "        return int(np.ceil(dim // self.n_blocks))\n",
    "    \n",
    "    def forward(self, idx):\n",
    "        # Hashed shape: (bs, seq_len, n_hashes)\n",
    "        xhash = self.process_fn(idx)\n",
    "        # Each embedding takes in a tensor of shape (bs, seq_len).\n",
    "        res_blocks = [e(hashed.squeeze()) for e, hashed in \n",
    "                      zip(self.emb, torch.chunk(xhash, xhash.shape[0], -1))]\n",
    "        return torch.cat(res_blocks, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-21T21:35:24.968586Z",
     "start_time": "2020-11-21T21:35:24.924031Z"
    }
   },
   "outputs": [],
   "source": [
    "def reduction_ratio(ax, vocab_size, emb_dim):\n",
    "    \"\"\"For testing purposes. Lets us compare the number of weights in a\n",
    "    traditional embedding matrix vs. the number of weights in our axial\n",
    "    encoding.\n",
    "    \"\"\"\n",
    "    normal_n = vocab_size * emb_dim\n",
    "    ax_n = sum(e.weight.numel() for e in ax.emb)\n",
    "    print('Normal embedding weights:', normal_n)\n",
    "    print('Axial encoding weights:', ax_n)\n",
    "    print('Difference:', normal_n - ax_n)\n",
    "    print('Ratio:', normal_n / ax_n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-21T21:35:25.164578Z",
     "start_time": "2020-11-21T21:35:25.118152Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([12, 2])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "AxialEncoding(\n",
       "  (emb): ModuleList(\n",
       "    (0): InitializedEmbedding(174, 50)\n",
       "    (1): InitializedEmbedding(174, 50)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab_size = 30_000\n",
    "emb_dim = 100\n",
    "bs = 12\n",
    "\n",
    "ax = AxialEncoding(vocab_size, emb_dim)\n",
    "x = torch.randint(0, vocab_size, (bs, 2))\n",
    "print(x.shape)\n",
    "ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([12, 2, 100])\n"
     ]
    }
   ],
   "source": [
    "res = ax(x)\n",
    "print(res.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Normal embedding weights: 3000000\n",
      "Axial encoding weights: 17400\n",
      "Difference: 2982600\n",
      "Ratio: 172.41379310344828\n"
     ]
    }
   ],
   "source": [
    "reduction_ratio(ax, vocab_size, emb_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([12, 2])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "MultiAxialEncoding(\n",
       "  (emb): ModuleList(\n",
       "    (0): Embedding(14, 25)\n",
       "    (1): Embedding(14, 25)\n",
       "    (2): Embedding(14, 25)\n",
       "    (3): Embedding(14, 25)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab_size = 30_000\n",
    "emb_dim = 100\n",
    "bs = 12\n",
    "\n",
    "ax = MultiAxialEncoding(vocab_size, emb_dim, 4)\n",
    "x = torch.randint(0, vocab_size, (bs, 2))\n",
    "print(x.shape)\n",
    "ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([12, 2, 100])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res1 = ax(x)\n",
    "res1.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MultiAxialEncoding(\n",
       "  (emb): ModuleList(\n",
       "    (0): Embedding(14, 25)\n",
       "    (1): Embedding(14, 25)\n",
       "    (2): Embedding(14, 25)\n",
       "    (3): Embedding(14, 25)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab_size = 30_000\n",
    "emb_dim = 100\n",
    "bs = 12\n",
    "\n",
    "ax_pre = MultiAxialEncoding(vocab_size, emb_dim, 4, pre_hashed=True)\n",
    "ax_pre"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By setting the weights of our pre-hashed embedding to the weights of our hashing embedding, we can check that the outputs are ultimately the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for e, e_pre in zip(ax.emb, ax_pre.emb):\n",
    "    e_pre.weight.data = e.weight.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([12, 2, 100])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xhash = probabilistic_hash_tensor(x, 14, 4)\n",
    "res2 = ax_pre(xhash)\n",
    "res2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(True)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(res1 == res2).all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Normal embedding weights: 3000000\n",
      "Axial encoding weights: 1400\n",
      "Difference: 2998600\n",
      "Ratio: 2142.8571428571427\n"
     ]
    }
   ],
   "source": [
    "reduction_ratio(ax_pre, vocab_size, emb_dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I imagine that as we increase `n_blocks`, there's likely a point where we simply won't have enough weights to encode the amount of information that's present in the data. It would take some experimentation to find where that line is, however."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MultiAxialEncoding(\n",
       "  (emb): ModuleList(\n",
       "    (0): Embedding(4, 12)\n",
       "    (1): Embedding(4, 12)\n",
       "    (2): Embedding(4, 12)\n",
       "    (3): Embedding(4, 12)\n",
       "    (4): Embedding(4, 12)\n",
       "    (5): Embedding(4, 12)\n",
       "    (6): Embedding(4, 12)\n",
       "    (7): Embedding(4, 12)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ax_large = MultiAxialEncoding(vocab_size, emb_dim, 8, pre_hashed=True)\n",
    "ax_large"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Normal embedding weights: 3000000\n",
      "Axial encoding weights: 384\n",
      "Difference: 2999616\n",
      "Ratio: 7812.5\n"
     ]
    }
   ],
   "source": [
    "reduction_ratio(ax_large, vocab_size, emb_dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Attention\n",
    "\n",
    "Some GPT2-esque layers. These were mostly intuition-building exercises - more thoroughly tested implementations likely exist in Huggingface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-01-03T04:16:36.635950Z",
     "start_time": "2021-01-03T04:16:36.550127Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class Projector(nn.Module):\n",
    "    \"\"\"Project input into multiple spaces. Used in DotProductAttention to\n",
    "    generate queries/keys/values.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, n_in, n_out_single=None, spaces=3):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        n_in: int\n",
    "            Size of input feature dimension, where input is (bs, n_in) or\n",
    "            (bs, seq_len, n_in). If the latter, this ONLY transforms the last \n",
    "            dimension. If you want to take multiple dimensions of information \n",
    "            into account simultaneously, you can flatten the input prior to \n",
    "            passing it in.\n",
    "        n_out_single: int or None\n",
    "            This determines the size of the feature dimension in each new \n",
    "            space. By default, this will be the same as n_in.\n",
    "        spaces: int\n",
    "            Number of spaces to project the input into. Default is 3 because\n",
    "            we commonly use this to generate queries, keys, and values for\n",
    "            attention computations.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.spaces = spaces\n",
    "        self.n_in = n_in\n",
    "        self.n_out_single = n_out_single or n_in\n",
    "        self.spaces = spaces\n",
    "        self.n_out = self.n_out_single * self.spaces\n",
    "        self.fc = nn.Linear(self.n_in, self.n_out)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        x: torch.Tensor\n",
    "            Shape (bs, n_in) or (bs, seq_len, n_in).\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        tuple[torch.Tensor]: Tuple of `spaces` tensors where each tensor has\n",
    "        shape (bs, n_out_single) or (bs, seq_len, n_out_single), depending on\n",
    "        the input shape.\n",
    "        \"\"\"\n",
    "        return self.fc(x).chunk(self.spaces, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class DotProductAttention(nn.Module):\n",
    "    \"\"\"GPT2-style attention block. This was mostly an intuition-building\n",
    "    exercise - in practice, Huggingface provides layers that should probably\n",
    "    be used instead.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, n_in, n_out=None, nf=None, n_heads=12,\n",
    "                 temperature='auto', p1=0.1, p2=0.1, return_attn=False):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        n_in: int\n",
    "            Last dimension of input, usually embedding dimension.\n",
    "        n_out: int or None\n",
    "            Size of output vectors. By default, this will be the same as the \n",
    "            input.\n",
    "        nf: int or None\n",
    "            Size (\"nf = number of features\") of queries/keys/values. \n",
    "            By default, this will be the same as n_in. Must be divisible by\n",
    "            n_heads.\n",
    "        n_heads: int\n",
    "            Number of attention heads to use. nf must be divisible\n",
    "            by this as each projected vector will be divided evenly among\n",
    "            each head.\n",
    "        temperature: str or float\n",
    "            If str, must be \"auto\", meaning softmax inputs will be scaled by\n",
    "            sqrt(n_proj_single). You can also specify a float, where values\n",
    "            <1 sharpen the distribution (usually not what we want here) and\n",
    "            values greater than one soften it (allowing attention head to \n",
    "            route more information from multiple neurons rather than almost \n",
    "            all from one).\n",
    "        p1: float\n",
    "            Value in (0.0, 1.0) setting the dropout probability on the \n",
    "            attention weights.\n",
    "        p2: float\n",
    "            Value in (0.0, 1.0) setting dropout probability following the \n",
    "            output layer.\n",
    "        return_attn: bool\n",
    "            If True, the `forward` method will return a tuple of \n",
    "            (output, attention_weights) tensors. If False (the default), just\n",
    "            return the output tensor.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        nf = nf or n_in\n",
    "        n_out = n_out or n_in\n",
    "        assert nf % n_heads == 0, \\\n",
    "            'n_proj_single must be divisible by n_heads'\n",
    "        \n",
    "        self.proj_in = Projector(n_in, nf, spaces=3)\n",
    "        # Reshape so hidden dimension is split equally between each head.\n",
    "        self.head_splitter = Rearrange('bs seq (heads f) -> bs heads seq f', \n",
    "                                       heads=n_heads)\n",
    "        self.soft = SmoothSoftmax(temperature)\n",
    "        self.drop_attn = nn.Dropout(p1)\n",
    "        # Concatenate output of each head.\n",
    "        self.head_merger = Rearrange('bs heads seq f -> bs seq (heads f)')\n",
    "        self.fc_out = nn.Linear(nf, n_out)\n",
    "        self.drop_out = nn.Dropout(p2)\n",
    "    \n",
    "        # Non-layer attributes.\n",
    "        self.n_heads = n_heads\n",
    "        self.temperature = temperature\n",
    "        self.p1 = p1\n",
    "        self.p2 = p2\n",
    "        self.return_attn = return_attn\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        x: torch.Tensor\n",
    "            Shape (bs, seq_len, n_in). n_in will usually be the sum of\n",
    "            embedding dimensions (word and positional). For other problems\n",
    "            (e.g. web browsing sequence classificaiton), this might include\n",
    "            other features about the page at time step T.\n",
    "        \"\"\"\n",
    "        q, k, v = map(self.head_splitter, self.proj_in(x))\n",
    "        scores = q @ k.transpose(-2, -1)\n",
    "        weights = self.drop_attn(self.soft(scores))\n",
    "        x = weights @ v\n",
    "        x = self.head_merger(x)\n",
    "        x = self.drop_out(self.fc_out(x))\n",
    "        return (x, weights) if self.return_attn else x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Bases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-19T04:24:57.756007Z",
     "start_time": "2020-10-19T04:24:57.711779Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class SiameseBase(BaseModel, ABC):\n",
    "    \"\"\"Parent class to implement a Siamese network or triplet network (or any\n",
    "    network that passes n inputs of the same shape through a shared encoder).\n",
    "    It concatenates the items into a single batch so the encoder's forward\n",
    "    method (implemented as self._forward) only needs to be called once.\n",
    "    \"\"\"\n",
    "\n",
    "    def forward(self, *xb):\n",
    "        bs = xb[0].shape[0]\n",
    "        xb = self._forward(torch.cat(xb, dim=0))\n",
    "        return xb.view(bs, -1, *xb.shape[1:])\n",
    "\n",
    "    @abstractmethod\n",
    "    def _forward(self, xb):\n",
    "        \"\"\"Forward pass for a single batch of x. Note that the batch dimension\n",
    "        here will be batch_size * n, where n is the number of images in a \n",
    "        single example (e.g. n=2 for a traditional Siamese Network, but you \n",
    "        can go arbitrarily high).\n",
    "        \"\"\"\n",
    "        raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-19T04:24:58.012054Z",
     "start_time": "2020-10-19T04:24:57.955297Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[torch.Size([4, 3, 8, 8]), torch.Size([4, 3, 8, 8]), torch.Size([4, 3, 8, 8])]"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs, c, h, w = 4, 3, 8, 8\n",
    "n = 3\n",
    "xb = [torch.randn(bs, c, h, w) for _ in range(n)]\n",
    "smap(*xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-19T04:24:58.438102Z",
     "start_time": "2020-10-19T04:24:58.400815Z"
    }
   },
   "outputs": [],
   "source": [
    "class TripletNet(SiameseBase):\n",
    "    \n",
    "    def __init__(self, c_in=3):\n",
    "        super().__init__()\n",
    "        self.conv = nn.Conv2d(c_in, 16, kernel_size=3, stride=2)\n",
    "        self.pool = nn.AdaptiveAvgPool2d((1, 1))\n",
    "        \n",
    "    def _forward(self, xb):\n",
    "        print(xb.shape)\n",
    "        xb = self.conv(xb)\n",
    "        print(xb.shape)\n",
    "        xb = self.pool(xb)\n",
    "        print(xb.shape)\n",
    "        xb = xb.squeeze(-1).squeeze(-1)\n",
    "        print(xb.shape)\n",
    "        return xb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example, each image is encoded as a 16D vector. We have 3 images per row and 4 rows per batch so we end up with a tensor of shape (4, 3, 16). Notice we only perform 1 forward pass: while we could simply define a separate encoder and pass each image through it separately (e.g. `[self.encoder(x) for x in xb]`), this becomes rather slow if n is large or if our encoder is enormous."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-19T04:24:59.454501Z",
     "start_time": "2020-10-19T04:24:59.415097Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([12, 3, 8, 8])\n",
      "torch.Size([12, 16, 3, 3])\n",
      "torch.Size([12, 16, 1, 1])\n",
      "torch.Size([12, 16])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 3, 16])"
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tnet = TripletNet()\n",
    "yh = tnet(*xb)\n",
    "yh.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our name TripletNet was slightly misleading here: the network can actually handle any choice of n. For instance, here we use it as a Siamese Net."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-19T04:25:49.234539Z",
     "start_time": "2020-10-19T04:25:49.188183Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([8, 3, 8, 8])\n",
      "torch.Size([8, 16, 3, 3])\n",
      "torch.Size([8, 16, 1, 1])\n",
      "torch.Size([8, 16])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 2, 16])"
      ]
     },
     "execution_count": 73,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "yh = tnet(*xb[:2])\n",
    "yh.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is often useful to extract intermediate activations from a model. We provide a convenient way to make a new model do this (or convert an existing Sequential model to do this)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class SequentialWithActivations(nn.Sequential):\n",
    "    \n",
    "    def __init__(self, *args, return_idx=()):\n",
    "        \"\"\"Create a sequential model that also returns activations from one or \n",
    "        more intermediate layers.         \n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        args: nn.Modules\n",
    "            Just like a Sequential model: pass in 1 or more layers in the \n",
    "            order you want them to process inputs.\n",
    "        return_idx: Iterable[int]\n",
    "            Indices of which layer outputs to return. Do not include the final\n",
    "            layer since that is always returned automatically. Activations \n",
    "            will be returned in increasing order by index - if you create a 4 \n",
    "            layer network and pass in return_idx=[2, 0], your output will \n",
    "            still be [layer_0_acts, layer_2_acts, final_layer_acts].\n",
    "            We recommend passing in indices in the expected return order to \n",
    "            avoid confusion.\n",
    "        \"\"\"\n",
    "        super().__init__(*args)\n",
    "        assert all(i < len(args) - 1 for i in return_idx), 'All ids in ' \\\n",
    "            'return_idx must correspond to layers before the final layer, ' \\\n",
    "            'which is always returned.'\n",
    "        self.return_idx = set(return_idx)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Returns\n",
    "        -------\n",
    "        Tuple[torch.Tensor]: N tensors where the first N-1 correspond to \n",
    "        self.return_idx (sorted in ascending order) and the last item is the \n",
    "        output of the final layer.\n",
    "        \"\"\"\n",
    "        res = []\n",
    "        for i, module in enumerate(self):\n",
    "            x = module(x)\n",
    "            if i in self.return_idx: res.append(x)\n",
    "        return (*res, x)\n",
    "        \n",
    "    @classmethod\n",
    "    def from_sequential(cls, model, return_idx=()):\n",
    "        \"\"\"Convert a standard Sequential model to a MultiOutputSequential.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        model: nn.Sequential\n",
    "        return_idx: Iterable[int]\n",
    "            Indices of which layer outputs to return. Do not include the final\n",
    "            layer since that is always returned automatically.\n",
    "        \"\"\"\n",
    "        model = copy.deepcopy(model)\n",
    "        return cls(*list(model), return_idx=return_idx)"
   ]
  }
 ],
 "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": 4
}
