Gaussian_Variational_Dropout.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.utils.data as data\n",
    "from torch.autograd import Variable\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "from tqdm import tqdm_notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_dataset(dataset='MNIST', dataset_dir='./data', batch_size=100):\n",
    "    dataset_ = {\n",
    "        'MNIST': datasets.MNIST,\n",
    "        'CIFAR10': datasets.CIFAR10\n",
    "    }[dataset]\n",
    "    \n",
    "    transform = {\n",
    "        'MNIST': transforms.ToTensor(),\n",
    "        'CIFAR10': transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "            ])\n",
    "    }[dataset]\n",
    "    \n",
    "    train_dataset = dataset_(root=dataset_dir,\n",
    "                             train=True,\n",
    "                             transform=transform,\n",
    "                             download=True)\n",
    "\n",
    "    train_loader = data.DataLoader(dataset=train_dataset,\n",
    "                                  batch_size=batch_size,\n",
    "                                  shuffle=True)\n",
    "\n",
    "    test_dataset = dataset_(root=dataset_dir,\n",
    "                             train=False,\n",
    "                             transform=transform,\n",
    "                             download=True)\n",
    "\n",
    "    test_loader = data.DataLoader(dataset=test_dataset,\n",
    "                                  batch_size=batch_size,\n",
    "                                  shuffle=False)\n",
    "    \n",
    "    return train_loader, test_loader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Standard Dropout\n",
    "\n",
    "$h_i = W_ix^{noise} = W^{noise}_ix = W_ix * noise_i = \\hat{h}_i * noise_i$\n",
    "- Masking input = Masking corresponding weight column\n",
    "- Masking activation = Masking corresponding weight row\n",
    "- Mask = Multiplicative Bernoulli random variable\n",
    "\n",
    "$h_i = \\hat{h}_i\\epsilon_i \\ /\\ (1-p)$\n",
    "- $\\epsilon_i \\sim Bernoulli(1-p)$\n",
    "- $p$: zero-out probability\n",
    "- Dividing by $1-p$ for train/test scale balance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gaussian Approximation of $\\epsilon\\ /\\ (1-p)$\n",
    "$\\epsilon \\sim Bernoulli(1-p)$\n",
    "- $\\mathbb{E}[\\epsilon]= 1-p$\n",
    "- $Var[\\epsilon] = \\frac{p}{1-p}$\n",
    "\n",
    "$z \\sim \\epsilon\\ /\\ (1-p)$\n",
    "- $\\mathbb{E}[z] = \\frac{(1-p)}{(1-p)} = 1$\n",
    "- $Var[z] = \\frac{p(1-p)}{(1-p)^2} = \\frac{p}{1-p}$\n",
    "\n",
    "Let's sample multiplicative noise directly from $N(1, \\frac{p}{1-p})$ instead of $\\frac{Bernoulli(1-p)}{1-p}$!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gaussian Dropout\n",
    "\n",
    "$h_i = \\hat{h}_i\\epsilon_i$\n",
    "- $\\epsilon_i \\sim N(1,\\alpha)$\n",
    "- $\\alpha: \\frac{p}{1-p}$\n",
    "- $p$: zero-out probability"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Variational Dropout\n",
    "\n",
    "- Variational Inference for Gaussian Dropout\n",
    "    - $D_{KL}[q_\\alpha(W)||p(w)] \\approx constant - 0.5\\log(\\alpha) - c_1\\alpha - c_2\\alpha - c_3\\alpha$\n",
    "    - $c_1 = 1.16145124, c_2 = −1.50204118, c_3 = 0.58629921$\n",
    "- Minimizing $D_{KL}$ (Adaptive regularizer)\n",
    "    - Train $\\alpha$ adaptively to data\n",
    "\n",
    "### Independent weight noise\n",
    "- $h_i = W_i^{noise}x$\n",
    "- $W^{noise}_{ij} = W_{ij}\\epsilon_{ij}$\n",
    "- $ \\epsilon_{ij} \\sim N(1,\\alpha)$\n",
    "\n",
    "### Correlated weight noise\n",
    "- $h_i = W_i^{noise}x$\n",
    "- $W^{noise}_{i} = W_{i}\\epsilon_{i}$\n",
    "- $ \\epsilon_{i} \\sim N(1,\\alpha)$\n",
    "\n",
    "### With Local Reparametrization Trick\n",
    "- $h_i = \\hat{h}_i\\epsilon_i$\n",
    "- $ \\epsilon_{i} \\sim N(1,\\alpha)$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GaussianDropout(nn.Module):\n",
    "    def __init__(self, alpha=1.0):\n",
    "        super(GaussianDropout, self).__init__()\n",
    "        self.alpha = torch.Tensor([alpha])\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Sample noise   e ~ N(1, alpha)\n",
    "        Multiply noise h = h_ * e\n",
    "        \"\"\"\n",
    "        if self.train():\n",
    "            # N(1, alpha)\n",
    "            epsilon = torch.randn(x.size()) * self.alpha + 1\n",
    "\n",
    "            epsilon = Variable(epsilon)\n",
    "            if x.is_cuda:\n",
    "                epsilon = epsilon.cuda()\n",
    "\n",
    "            return x * epsilon\n",
    "        else:\n",
    "            return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VariationalDropout(nn.Module):\n",
    "    def __init__(self, alpha=1.0, dim=None):\n",
    "        super(VariationalDropout, self).__init__()\n",
    "        \n",
    "        self.dim = dim\n",
    "        self.max_alpha = alpha\n",
    "        # Initial alpha\n",
    "        log_alpha = (torch.ones(dim) * alpha).log()\n",
    "        self.log_alpha = nn.Parameter(log_alpha)\n",
    "        \n",
    "    def kl(self):\n",
    "        c1 = 1.16145124\n",
    "        c2 = -1.50204118\n",
    "        c3 = 0.58629921\n",
    "        \n",
    "        alpha = self.log_alpha.exp()\n",
    "        \n",
    "        negative_kl = 0.5 * self.log_alpha + c1 * alpha + c2 * alpha**2 + c3 * alpha**3\n",
    "        \n",
    "        kl = -negative_kl\n",
    "        \n",
    "        return kl.mean()\n",
    "    \n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Sample noise   e ~ N(1, alpha)\n",
    "        Multiply noise h = h_ * e\n",
    "        \"\"\"\n",
    "        if self.train():\n",
    "            # N(0,1)\n",
    "            epsilon = Variable(torch.randn(x.size()))\n",
    "            if x.is_cuda:\n",
    "                epsilon = epsilon.cuda()\n",
    "\n",
    "            # Clip alpha\n",
    "            self.log_alpha.data = torch.clamp(self.log_alpha.data, max=self.max_alpha)\n",
    "            alpha = self.log_alpha.exp()\n",
    "\n",
    "            # N(1, alpha)\n",
    "            epsilon = epsilon * alpha\n",
    "\n",
    "            return x * epsilon\n",
    "        else:\n",
    "            return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dropout(p=None, dim=None, method='standard'):\n",
    "    if method == 'standard':\n",
    "        return nn.Dropout(p)\n",
    "    elif method == 'gaussian':\n",
    "        return GaussianDropout(p/(1-p))\n",
    "    elif method == 'variational':\n",
    "        return VariationalDropout(p/(1-p), dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self,\n",
    "                 image_dim=28*28,\n",
    "                 dropout_method='standard'):\n",
    "        super(Net, self).__init__()\n",
    "        \"\"\"3-Layer Fully-connected NN\"\"\"\n",
    "        \n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(image_dim, 500),\n",
    "            dropout(0.2, 500, dropout_method),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(500, 100),\n",
    "            dropout(0.5, 100, dropout_method),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(100, 10)\n",
    "        )\n",
    "        \n",
    "    def kl(self):\n",
    "        kl = 0\n",
    "        for name, module in self.net.named_modules():\n",
    "            if isinstance(module, VariationalDropout):\n",
    "                kl += module.kl().sum()\n",
    "        return kl\n",
    "        \n",
    "            \n",
    "    def forward(self, x):\n",
    "        return self.net(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Solver(object):\n",
    "    def __init__(self, dropout_method='standard', dataset='MNIST', n_epochs=50, lr=0.005):\n",
    "        self.n_epochs = n_epochs\n",
    "        \n",
    "        self.train_loader, self.test_loader = build_dataset(dataset, './data')\n",
    "        \n",
    "        self.image_dim = {'MNIST': 28*28, 'CIFAR10': 3*32*32}[dataset]\n",
    "        \n",
    "        self.dropout_method = dropout_method\n",
    "        \n",
    "        self.net = Net(\n",
    "            image_dim=self.image_dim,\n",
    "            dropout_method=dropout_method).cuda()\n",
    "        self.loss_fn = nn.CrossEntropyLoss()\n",
    "        self.optimizer = optim.Adam(self.net.parameters(), lr=lr)\n",
    "                \n",
    "    def train(self):\n",
    "        self.net.train()\n",
    "        \n",
    "        for epoch_i in tqdm_notebook(range(self.n_epochs)):\n",
    "            epoch_i += 1\n",
    "            epoch_loss = 0\n",
    "            epoch_kl = 0\n",
    "            for images, labels in self.train_loader:\n",
    "                images = Variable(images).view(-1, self.image_dim).cuda()\n",
    "                labels = Variable(labels).cuda()\n",
    "\n",
    "                logits = self.net(images)\n",
    "                \n",
    "                loss = self.loss_fn(logits, labels)\n",
    "                \n",
    "                if self.dropout_method == 'variational':\n",
    "                    kl = self.net.kl()\n",
    "                    total_loss = loss + kl / 10\n",
    "                else:\n",
    "                    total_loss = loss\n",
    "\n",
    "                self.optimizer.zero_grad()\n",
    "                total_loss.backward()\n",
    "\n",
    "                self.optimizer.step()\n",
    "                \n",
    "                epoch_loss += float(loss.data)\n",
    "                if self.dropout_method == 'variational':\n",
    "                    epoch_kl += float(kl.data)\n",
    "\n",
    "            if not self.dropout_method == 'variational':\n",
    "                epoch_loss /= len(self.train_loader.dataset)\n",
    "                print(f'Epoch {epoch_i} | loss: {epoch_loss:.4f}')\n",
    "\n",
    "            else:\n",
    "                epoch_loss /= len(self.train_loader.dataset)\n",
    "                epoch_kl /= len(self.train_loader.dataset)\n",
    "                print(f'Epoch {epoch_i} | loss: {epoch_loss:.4f}, kl: {epoch_kl:.4f}')\n",
    "            \n",
    "    def evaluate(self):\n",
    "        total = 0\n",
    "        correct = 0\n",
    "        self.net.eval()\n",
    "        for images, labels in self.test_loader:\n",
    "            images = Variable(images).view(-1, self.image_dim).cuda()\n",
    "\n",
    "            logits = self.net(images)\n",
    "            \n",
    "            _, predicted = torch.max(logits.data, 1)\n",
    "            \n",
    "            total += labels.size(0)\n",
    "            correct += (predicted.cpu() == labels).sum()\n",
    "                \n",
    "        print(f'Accuracy: {100 * correct / total:.2f}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Experiments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "standard_solver = Solver('standard')\n",
    "\n",
    "standard_solver.train()\n",
    "standard_solver.evaluate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "gaussian_solver = Solver('gaussian')\n",
    "\n",
    "gaussian_solver.train()\n",
    "gaussian_solver.evaluate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "variational_solver = Solver('variational')\n",
    "\n",
    "variational_solver.train()\n",
    "variational_solver.evaluate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:mldemo]",
   "language": "python",
   "name": "conda-env-mldemo-py"
  },
  "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.6.3"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}