batch-norm/Batch_Normalization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Batch Normalization \n",
    "\n",
    "Batch normalization was introduced in Sergey Ioffe's and Christian Szegedy's 2015 paper [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://arxiv.org/pdf/1502.03167.pdf). The idea is that, instead of just normalizing the inputs to the network, we normalize the inputs to _layers within_ the network. \n",
    "> It's called **batch** normalization because during training, we normalize each layer's inputs by using the mean and variance of the values in the current *batch*."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Batch Normalization in PyTorch<a id=\"implementation_1\"></a>\n",
    "\n",
    "This section of the notebook shows you one way to add batch normalization to a neural network built in PyTorch. \n",
    "\n",
    "The following cells import the packages we need in the notebook and load the MNIST dataset to use in our experiments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from torchvision import datasets\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "# number of subprocesses to use for data loading\n",
    "num_workers = 0\n",
    "# how many samples per batch to load\n",
    "batch_size = 64\n",
    "\n",
    "# convert data to torch.FloatTensor\n",
    "transform = transforms.ToTensor()\n",
    "\n",
    "# get the training and test datasets\n",
    "train_data = datasets.MNIST(root='data', train=True,\n",
    "                            download=True, transform=transform)\n",
    "\n",
    "test_data = datasets.MNIST(root='data', train=False,\n",
    "                           download=True, transform=transform)\n",
    "\n",
    "# prepare data loaders\n",
    "train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size,\n",
    "                                           num_workers=num_workers)\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size,\n",
    "                                          num_workers=num_workers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# obtain one batch of training images\n",
    "dataiter = iter(train_loader)\n",
    "images, labels = dataiter.next()\n",
    "images = images.numpy()\n",
    "\n",
    "# get one image from the batch\n",
    "img = np.squeeze(images[0])\n",
    "\n",
    "fig = plt.figure(figsize = (3,3)) \n",
    "ax = fig.add_subplot(111)\n",
    "ax.imshow(img, cmap='gray')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Neural network classes for testing\n",
    "\n",
    "The following class, `NeuralNet`, allows us to create identical neural networks **with and without batch normalization** to compare. The code is heavily documented, but there is also some additional discussion later. You do not need to read through it all before going through the rest of the notebook, but the comments within the code blocks may answer some of your questions.\n",
    "\n",
    "*About the code:*\n",
    ">We are defining a simple MLP for classification; this design choice was made to support the discussion related to batch normalization and not to get the best classification accuracy.\n",
    "\n",
    "### (Important) Model Details\n",
    "\n",
    "There are quite a few comments in the code, so those should answer most of your questions. However, let's take a look at the most important lines.\n",
    "\n",
    "We add batch normalization to layers inside the `__init__` function. Here are some important points about that code:\n",
    "1. Layers with batch normalization do **not** include a bias term.\n",
    "2. We use PyTorch's [BatchNorm1d](https://pytorch.org/docs/stable/nn.html#batchnorm1d) function to handle the math. This is the function you use to operate on linear layer outputs; you'll use [BatchNorm2d](https://pytorch.org/docs/stable/nn.html#batchnorm2d) for 2D outputs like filtered images from convolutional layers. \n",
    "3. We add the batch normalization layer **before** calling the activation function.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class NeuralNet(nn.Module):\n",
    "    def __init__(self, use_batch_norm, input_size=784, hidden_dim=256, output_size=10):\n",
    "        \"\"\"\n",
    "        Creates a PyTorch net using the given parameters.\n",
    "        \n",
    "        :param use_batch_norm: bool\n",
    "            Pass True to create a network that uses batch normalization; False otherwise\n",
    "            Note: this network will not use batch normalization on layers that do not have an\n",
    "            activation function.\n",
    "        \"\"\"\n",
    "        super(NeuralNet, self).__init__() # init super\n",
    "        \n",
    "        # Default layer sizes\n",
    "        self.input_size = input_size # (28*28 images)\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.output_size = output_size # (number of classes)\n",
    "        # Keep track of whether or not this network uses batch normalization.\n",
    "        self.use_batch_norm = use_batch_norm\n",
    "        \n",
    "        # define hidden linear layers, with optional batch norm on their outputs\n",
    "        # layers with batch_norm applied have no bias term\n",
    "        if use_batch_norm:\n",
    "            self.fc1 = nn.Linear(input_size, hidden_dim*2, bias=False)\n",
    "            self.batch_norm1 = nn.BatchNorm1d(hidden_dim*2)\n",
    "        else:\n",
    "            self.fc1 = nn.Linear(input_size, hidden_dim*2)\n",
    "            \n",
    "        # define *second* hidden linear layers, with optional batch norm on their outputs\n",
    "        if use_batch_norm:\n",
    "            self.fc2 = nn.Linear(hidden_dim*2, hidden_dim, bias=False)\n",
    "            self.batch_norm2 = nn.BatchNorm1d(hidden_dim)\n",
    "        else:\n",
    "            self.fc2 = nn.Linear(hidden_dim*2, hidden_dim)\n",
    "        \n",
    "        # third and final, fully-connected layer\n",
    "        self.fc3 = nn.Linear(hidden_dim, output_size)\n",
    "        \n",
    "        \n",
    "    def forward(self, x):\n",
    "        # flatten image\n",
    "        x = x.view(-1, 28*28)\n",
    "        # all hidden layers + optional batch norm + relu activation\n",
    "        x = self.fc1(x)\n",
    "        if self.use_batch_norm:\n",
    "            x = self.batch_norm1(x)\n",
    "        x = F.relu(x)\n",
    "        # second layer\n",
    "        x = self.fc2(x)\n",
    "        if self.use_batch_norm:\n",
    "            x = self.batch_norm2(x)\n",
    "        x = F.relu(x)\n",
    "        # third layer, no batch norm or activation\n",
    "        x = self.fc3(x)\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create two different models for testing\n",
    "\n",
    "* `net_batchnorm` is a linear classification model **with** batch normalization applied to the output of its hidden layers\n",
    "* `net_no_norm` is a plain MLP, without batch normalization\n",
    "\n",
    "Besides the normalization layers, everthing about these models is the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net_batchnorm = NeuralNet(use_batch_norm=True)\n",
    "net_no_norm = NeuralNet(use_batch_norm=False)\n",
    "\n",
    "print(net_batchnorm)\n",
    "print()\n",
    "print(net_no_norm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "---\n",
    "## Training\n",
    "\n",
    "The below `train` function will take in a model and some number of epochs. We'll use cross entropy loss and stochastic gradient descent for optimization. This function returns the losses, recorded after each epoch, so that we can display and compare the behavior of different models.\n",
    "\n",
    "#### `.train()` mode\n",
    "Note that we tell our model whether or not it should be in training mode, `model.train()`. This is an important step because batch normalization has different behavior during training on a batch or testing/evaluating on a larger dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def train(model, n_epochs=10):\n",
    "    # number of epochs to train the model\n",
    "    n_epochs = n_epochs\n",
    "    # track losses\n",
    "    losses = []\n",
    "        \n",
    "    # optimization strategy \n",
    "    # specify loss function (categorical cross-entropy)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "    # specify optimizer (stochastic gradient descent) and learning rate = 0.01\n",
    "    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)\n",
    "\n",
    "    # set the model to training mode\n",
    "    model.train()\n",
    "    \n",
    "    for epoch in range(1, n_epochs+1):\n",
    "        # monitor training loss\n",
    "        train_loss = 0.0\n",
    "\n",
    "        ###################\n",
    "        # train the model #\n",
    "        ###################\n",
    "        batch_count = 0\n",
    "        for batch_idx, (data, target) in enumerate(train_loader):\n",
    "            # clear the gradients of all optimized variables\n",
    "            optimizer.zero_grad()\n",
    "            # forward pass: compute predicted outputs by passing inputs to the model\n",
    "            output = model(data)\n",
    "            # calculate the loss\n",
    "            loss = criterion(output, target)\n",
    "            # backward pass: compute gradient of the loss with respect to model parameters\n",
    "            loss.backward()\n",
    "            # perform a single optimization step (parameter update)\n",
    "            optimizer.step()\n",
    "            # update average training loss\n",
    "            train_loss += loss.item() # add up avg batch loss\n",
    "            batch_count +=1                \n",
    "\n",
    "        # print training statistics \n",
    "        losses.append(train_loss/batch_count)\n",
    "        print('Epoch: {} \\tTraining Loss: {:.6f}'.format(\n",
    "            epoch, \n",
    "            train_loss/batch_count))\n",
    "    \n",
    "    # return all recorded batch losses\n",
    "    return losses\n",
    "        \n",
    "      "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparing Models\n",
    "\n",
    "In the below cells, we train our two different models and compare their trainining loss over time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# batchnorm model losses\n",
    "# this may take some time to train\n",
    "losses_batchnorm = train(net_batchnorm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# *no* norm model losses\n",
    "# you should already start to see a difference in training losses\n",
    "losses_no_norm = train(net_no_norm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compare\n",
    "fig, ax = plt.subplots(figsize=(12,8))\n",
    "#losses_batchnorm = np.array(losses_batchnorm)\n",
    "#losses_no_norm = np.array(losses_no_norm)\n",
    "plt.plot(losses_batchnorm, label='Using batchnorm', alpha=0.5)\n",
    "plt.plot(losses_no_norm, label='No norm', alpha=0.5)\n",
    "plt.title(\"Training Losses\")\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "---\n",
    "## Testing\n",
    "\n",
    "You should see that the model with batch normalization, starts off with a lower training loss and, over ten epochs of training, gets to a training loss that is noticeably lower than our model without normalization.\n",
    "\n",
    "Next, let's see how both these models perform on our test data! Below, we have a function `test` that takes in a model and a parameter `train` (True or False) which indicates whether the model should be in training or evaulation mode. This is for comparison purposes, later. This function will calculate some test statistics including the overall test accuracy of a passed in model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def test(model, train):\n",
    "    # initialize vars to monitor test loss and accuracy\n",
    "    class_correct = list(0. for i in range(10))\n",
    "    class_total = list(0. for i in range(10))\n",
    "    test_loss = 0.0\n",
    "\n",
    "    # set model to train or evaluation mode\n",
    "    # just to see the difference in behavior\n",
    "    if(train==True):\n",
    "        model.train()\n",
    "    if(train==False):\n",
    "        model.eval()\n",
    "    \n",
    "    # loss criterion\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    \n",
    "    for batch_idx, (data, target) in enumerate(test_loader):\n",
    "        batch_size = data.size(0)\n",
    "        # forward pass: compute predicted outputs by passing inputs to the model\n",
    "        output = model(data)\n",
    "        # calculate the loss\n",
    "        loss = criterion(output, target)\n",
    "        # update average test loss \n",
    "        test_loss += loss.item()*batch_size\n",
    "        # convert output probabilities to predicted class\n",
    "        _, pred = torch.max(output, 1)\n",
    "        # compare predictions to true label\n",
    "        correct = np.squeeze(pred.eq(target.data.view_as(pred)))\n",
    "        # calculate test accuracy for each object class\n",
    "        for i in range(batch_size):\n",
    "            label = target.data[i]\n",
    "            class_correct[label] += correct[i].item()\n",
    "            class_total[label] += 1\n",
    "\n",
    "    print('Test Loss: {:.6f}\\n'.format(test_loss/len(test_loader.dataset)))\n",
    "\n",
    "    for i in range(10):\n",
    "        if class_total[i] > 0:\n",
    "            print('Test Accuracy of %5s: %2d%% (%2d/%2d)' % (\n",
    "                str(i), 100 * class_correct[i] / class_total[i],\n",
    "                np.sum(class_correct[i]), np.sum(class_total[i])))\n",
    "        else:\n",
    "            print('Test Accuracy of %5s: N/A (no training examples)' % (classes[i]))\n",
    "\n",
    "    print('\\nTest Accuracy (Overall): %2d%% (%2d/%2d)' % (\n",
    "        100. * np.sum(class_correct) / np.sum(class_total),\n",
    "        np.sum(class_correct), np.sum(class_total)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training and Evaluation Mode\n",
    "\n",
    "Setting a model to evaluation mode is important for models with batch normalization layers!\n",
    "\n",
    ">* Training mode means that the batch normalization layers will use **batch** statistics to calculate the batch norm. \n",
    "* Evaluation mode, on the other hand, uses the estimated **population** mean and variance from the entire training set, which should give us increased performance on this test data! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test batchnorm case, in *train* mode\n",
    "test(net_batchnorm, train=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test batchnorm case, in *evaluation* mode\n",
    "test(net_batchnorm, train=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for posterity, test no norm case in eval mode\n",
    "test(net_no_norm, train=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Which model has the highest accuracy?\n",
    "\n",
    "You should see a small improvement when comparing the batch norm model's accuracy in training and evaluation mode; **evaluation mode** should give a small improvement!\n",
    "\n",
    "You should also see that the model that uses batch norm layers shows a marked improvement in overall accuracy when compared with the no-normalization model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Considerations for other network types\n",
    "\n",
    "This notebook demonstrates batch normalization in a standard neural network with fully connected layers. You can also use batch normalization in other types of networks, but there are some special considerations.\n",
    "\n",
    "### ConvNets\n",
    "\n",
    "Convolution layers consist of multiple feature maps. (Remember, the depth of a convolutional layer refers to its number of feature maps.) And the weights for each feature map are shared across all the inputs that feed into the layer. Because of these differences, batch normalizing convolutional layers requires batch/population mean and variance per feature map rather than per node in the layer.\n",
    "\n",
    "> To apply batch normalization on the outputs of convolutional layers, we use [BatchNorm2d](https://pytorch.org/docs/stable/nn.html#batchnorm2d)\n",
    "\n",
    "\n",
    "### RNNs\n",
    "\n",
    "Batch normalization can work with recurrent neural networks, too, as shown in the 2016 paper [Recurrent Batch Normalization](https://arxiv.org/abs/1603.09025). It's a bit more work to implement, but basically involves calculating the means and variances per time step instead of per layer. You can find an example where someone implemented recurrent batch normalization in PyTorch, in [this GitHub repo](https://github.com/jihunchoi/recurrent-batch-normalization-pytorch)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [default]",
   "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}