03_consistent_mc_dropout.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp consistent_mc_dropout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Appended /home/blackhc/PycharmProjects/blackhc.batchbald/src to paths\n",
      "Switched to directory /home/blackhc/PycharmProjects/blackhc.batchbald\n",
      "%load_ext autoreload\n",
      "%autoreload 2\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "import blackhc.project.script\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Consistent MC Dropout\n",
    "> Custom consistent dropout modules"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For BNNs, we are going to use MC dropout.\n",
    "\n",
    "To be able to compute BatchBALD scores, we need consistent MC dropout, which uses the consistent masks for inference. That means, that we draw $K$ masks and then keep them fixed while drawing the $K$ inference samples for each input in the test set.\n",
    "\n",
    "During training, masks are redrawn for every sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "import torch\n",
    "from torch.nn import Module"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bayesian Module\n",
    "\n",
    "To make this work in an efficient way, we are going to define an abstract wrapper module that takes a batch `input_B` and outputs `results_B_K`.\n",
    "\n",
    "Internally, it will blow up the input batch to $(B \\cdot K) \\times \\cdots$ and then pass it to `mc_forward_impl`, which should be overriden.\n",
    "\n",
    "`ConsistentMCDropout` layers will know to reshape the inputs to $B \\times K \\times \\cdots$ and apply consistent masks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# exports\n",
    "\n",
    "\n",
    "class BayesianModule(Module):\n",
    "    \"\"\"A module that we can sample multiple times from given a single input batch.\n",
    "\n",
    "    To be efficient, the module allows for a part of the forward pass to be deterministic.\n",
    "    \"\"\"\n",
    "\n",
    "    k = None\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    # Returns B x n x output\n",
    "    def forward(self, input_B: torch.Tensor, k: int):\n",
    "        BayesianModule.k = k\n",
    "\n",
    "        mc_input_BK = BayesianModule.mc_tensor(input_B, k)\n",
    "        mc_output_BK = self.mc_forward_impl(mc_input_BK)\n",
    "        mc_output_B_K = BayesianModule.unflatten_tensor(mc_output_BK, k)\n",
    "        return mc_output_B_K\n",
    "\n",
    "    def mc_forward_impl(self, mc_input_BK: torch.Tensor):\n",
    "        return mc_input_BK\n",
    "\n",
    "    @staticmethod\n",
    "    def unflatten_tensor(input: torch.Tensor, k: int):\n",
    "        input = input.view([-1, k] + list(input.shape[1:]))\n",
    "        return input\n",
    "\n",
    "    @staticmethod\n",
    "    def flatten_tensor(mc_input: torch.Tensor):\n",
    "        return mc_input.flatten(0, 1)\n",
    "\n",
    "    @staticmethod\n",
    "    def mc_tensor(input: torch.tensor, k: int):\n",
    "        mc_shape = [input.shape[0], k] + list(input.shape[1:])\n",
    "        return input.unsqueeze(1).expand(mc_shape).flatten(0, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Consistent MC Dropout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# exports\n",
    "\n",
    "\n",
    "class _ConsistentMCDropout(Module):\n",
    "    def __init__(self, p=0.5):\n",
    "        super().__init__()\n",
    "\n",
    "        if p < 0 or p > 1:\n",
    "            raise ValueError(\"dropout probability has to be between 0 and 1, \" \"but got {}\".format(p))\n",
    "\n",
    "        self.p = p\n",
    "        self.mask = None\n",
    "\n",
    "    def extra_repr(self):\n",
    "        return \"p={}\".format(self.p)\n",
    "\n",
    "    def reset_mask(self):\n",
    "        self.mask = None\n",
    "\n",
    "    def train(self, mode=True):\n",
    "        super().train(mode)\n",
    "        if not mode:\n",
    "            self.reset_mask()\n",
    "\n",
    "    def _get_sample_mask_shape(self, sample_shape):\n",
    "        return sample_shape\n",
    "\n",
    "    def _create_mask(self, input, k):\n",
    "        mask_shape = [1, k] + list(self._get_sample_mask_shape(input.shape[1:]))\n",
    "        mask = torch.empty(mask_shape, dtype=torch.bool, device=input.device).bernoulli_(self.p)\n",
    "        return mask\n",
    "\n",
    "    def forward(self, input: torch.Tensor):\n",
    "        if self.p == 0.0:\n",
    "            return input\n",
    "\n",
    "        k = BayesianModule.k\n",
    "        if self.training:\n",
    "            # Create a new mask on each call and for each batch element.\n",
    "            k = input.shape[0]\n",
    "            mask = self._create_mask(input, k)\n",
    "        else:\n",
    "            if self.mask is None:\n",
    "                # print('recreating mask', self)\n",
    "                # Recreate mask.\n",
    "                self.mask = self._create_mask(input, k)\n",
    "\n",
    "            mask = self.mask\n",
    "\n",
    "        mc_input = BayesianModule.unflatten_tensor(input, k)\n",
    "        mc_output = mc_input.masked_fill(mask, 0) / (1 - self.p)\n",
    "\n",
    "        # Flatten MCDI, batch into one dimension again.\n",
    "        return BayesianModule.flatten_tensor(mc_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "\n",
    "class ConsistentMCDropout(_ConsistentMCDropout):\n",
    "    r\"\"\"Randomly zeroes some of the elements of the input\n",
    "    tensor with probability :attr:`p` using samples from a Bernoulli\n",
    "    distribution. The elements to zero are randomized on every forward call during training time.\n",
    "\n",
    "    During eval time, a fixed mask is picked and kept until `reset_mask()` is called.\n",
    "\n",
    "    This has proven to be an effective technique for regularization and\n",
    "    preventing the co-adaptation of neurons as described in the paper\n",
    "    `Improving neural networks by preventing co-adaptation of feature\n",
    "    detectors`_ .\n",
    "\n",
    "    Furthermore, the outputs are scaled by a factor of :math:`\\frac{1}{1-p}` during\n",
    "    training. This means that during evaluation the module simply computes an\n",
    "    identity function.\n",
    "\n",
    "    Args:\n",
    "        p: probability of an element to be zeroed. Default: 0.5\n",
    "        inplace: If set to ``True``, will do this operation in-place. Default: ``False``\n",
    "\n",
    "    Shape:\n",
    "        - Input: `Any`. Input can be of any shape\n",
    "        - Output: `Same`. Output is of the same shape as input\n",
    "\n",
    "    Examples::\n",
    "\n",
    "        >>> m = nn.Dropout(p=0.2)\n",
    "        >>> input = torch.randn(20, 16)\n",
    "        >>> output = m(input)\n",
    "\n",
    "    .. _Improving neural networks by preventing co-adaptation of feature\n",
    "        detectors: https://arxiv.org/abs/1207.0580\n",
    "    \"\"\"\n",
    "    pass\n",
    "\n",
    "\n",
    "class ConsistentMCDropout2d(_ConsistentMCDropout):\n",
    "    r\"\"\"Randomly zeroes whole channels of the input tensor.\n",
    "    The channels to zero-out are randomized on every forward call.\n",
    "\n",
    "    During eval time, a fixed mask is picked and kept until `reset_mask()` is called.\n",
    "\n",
    "    Usually the input comes from :class:`nn.Conv2d` modules.\n",
    "\n",
    "    As described in the paper\n",
    "    `Efficient Object Localization Using Convolutional Networks`_ ,\n",
    "    if adjacent pixels within feature maps are strongly correlated\n",
    "    (as is normally the case in early convolution layers) then i.i.d. dropout\n",
    "    will not regularize the activations and will otherwise just result\n",
    "    in an effective learning rate decrease.\n",
    "\n",
    "    In this case, :func:`nn.Dropout2d` will help promote independence between\n",
    "    feature maps and should be used instead.\n",
    "\n",
    "    Args:\n",
    "        p (float, optional): probability of an element to be zero-ed.\n",
    "        inplace (bool, optional): If set to ``True``, will do this operation\n",
    "            in-place\n",
    "\n",
    "    Shape:\n",
    "        - Input: :math:`(N, C, H, W)`\n",
    "        - Output: :math:`(N, C, H, W)` (same shape as input)\n",
    "\n",
    "    Examples::\n",
    "\n",
    "        >>> m = nn.Dropout2d(p=0.2)\n",
    "        >>> input = torch.randn(20, 16, 32, 32)\n",
    "        >>> output = m(input)\n",
    "\n",
    "    .. _Efficient Object Localization Using Convolutional Networks:\n",
    "       http://arxiv.org/abs/1411.4280\n",
    "    \"\"\"\n",
    "\n",
    "    def _get_sample_mask_shape(self, sample_shape):\n",
    "        return [sample_shape[0]] + [1] * (len(sample_shape) - 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example\n",
    "\n",
    "The following defines a DNN module that can learn MNIST."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BayesianCNN(\n",
       "  (conv1): Conv2d(1, 32, kernel_size=(5, 5), stride=(1, 1))\n",
       "  (conv1_drop): ConsistentMCDropout2d(p=0.5)\n",
       "  (conv2): Conv2d(32, 64, kernel_size=(5, 5), stride=(1, 1))\n",
       "  (conv2_drop): ConsistentMCDropout2d(p=0.5)\n",
       "  (fc1): Linear(in_features=1024, out_features=128, bias=True)\n",
       "  (fc1_drop): ConsistentMCDropout(p=0.5)\n",
       "  (fc2): Linear(in_features=128, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from torch import nn as nn\n",
    "from torch.nn import functional as F\n",
    "\n",
    "\n",
    "class BayesianCNN(BayesianModule):\n",
    "    def __init__(self, num_classes=10):\n",
    "        super().__init__()\n",
    "\n",
    "        self.conv1 = nn.Conv2d(1, 32, kernel_size=5)\n",
    "        self.conv1_drop = ConsistentMCDropout2d()\n",
    "        self.conv2 = nn.Conv2d(32, 64, kernel_size=5)\n",
    "        self.conv2_drop = ConsistentMCDropout2d()\n",
    "        self.fc1 = nn.Linear(1024, 128)\n",
    "        self.fc1_drop = ConsistentMCDropout()\n",
    "        self.fc2 = nn.Linear(128, num_classes)\n",
    "\n",
    "    def mc_forward_impl(self, input: torch.Tensor):\n",
    "        input = F.relu(F.max_pool2d(self.conv1_drop(self.conv1(input)), 2))\n",
    "        input = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(input)), 2))\n",
    "        input = input.view(-1, 1024)\n",
    "        input = F.relu(self.fc1_drop(self.fc1(input)))\n",
    "        input = self.fc2(input)\n",
    "        input = F.log_softmax(input, dim=1)\n",
    "\n",
    "        return input\n",
    "\n",
    "\n",
    "BayesianCNN()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}