{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Early Stopping Example\n",
    "In this notebook, we will train an Multi-Layer Perceptron (MLP) to classify images from the [MNIST database](http://yann.lecun.com/exdb/mnist/) hand-written digit database, and use early stopping to stop the training when the model starts to overfit to the training data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import libraries\n",
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load and Batch the Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision import datasets\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data.sampler import SubsetRandomSampler\n",
    "\n",
    "def create_datasets(batch_size):\n",
    "\n",
    "    # percentage of training set to use as validation\n",
    "    valid_size = 0.2\n",
    "\n",
    "    # convert data to torch.FloatTensor\n",
    "    transform = transforms.ToTensor()\n",
    "\n",
    "    # choose the training and test datasets\n",
    "    train_data = datasets.MNIST(root='data', \n",
    "                                train=True,\n",
    "                                download=True, \n",
    "                                transform=transform)\n",
    "\n",
    "    test_data = datasets.MNIST(root='data',\n",
    "                               train=False,\n",
    "                               download=True,\n",
    "                               transform=transform)\n",
    "\n",
    "    # obtain training indices that will be used for validation\n",
    "    num_train = len(train_data)\n",
    "    indices = list(range(num_train))\n",
    "    np.random.shuffle(indices)\n",
    "    split = int(np.floor(valid_size * num_train))\n",
    "    train_idx, valid_idx = indices[split:], indices[:split]\n",
    "    \n",
    "    # define samplers for obtaining training and validation batches\n",
    "    train_sampler = SubsetRandomSampler(train_idx)\n",
    "    valid_sampler = SubsetRandomSampler(valid_idx)\n",
    "    \n",
    "    # load training data in batches\n",
    "    train_loader = torch.utils.data.DataLoader(train_data,\n",
    "                                               batch_size=batch_size,\n",
    "                                               sampler=train_sampler,\n",
    "                                               num_workers=0)\n",
    "    \n",
    "    # load validation data in batches\n",
    "    valid_loader = torch.utils.data.DataLoader(train_data,\n",
    "                                               batch_size=batch_size,\n",
    "                                               sampler=valid_sampler,\n",
    "                                               num_workers=0)\n",
    "    \n",
    "    # load test data in batches\n",
    "    test_loader = torch.utils.data.DataLoader(test_data,\n",
    "                                              batch_size=batch_size,\n",
    "                                              num_workers=0)\n",
    "    \n",
    "    return train_loader, test_loader, valid_loader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the Network\n",
    "Defining a simple MLP model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Net(\n",
      "  (fc1): Linear(in_features=784, out_features=128, bias=True)\n",
      "  (fc2): Linear(in_features=128, out_features=128, bias=True)\n",
      "  (fc3): Linear(in_features=128, out_features=10, bias=True)\n",
      "  (dropout): Dropout(p=0.5)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.fc1 = nn.Linear(28 * 28, 128)\n",
    "        self.fc2 = nn.Linear(128, 128)\n",
    "        self.fc3 = nn.Linear(128, 10)\n",
    "        self.dropout = nn.Dropout(0.5)\n",
    "    def forward(self, x):\n",
    "        # flatten image input\n",
    "        x = x.view(-1, 28 * 28)\n",
    "        # add hidden layer, with relu activation function\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = self.dropout(x)\n",
    "        # add hidden layer, with relu activation function\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.dropout(x)\n",
    "        # add output layer\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "\n",
    "# initialize the NN\n",
    "model = Net()\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Specify Loss Function and Optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# specify loss function\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "# specify optimizer\n",
    "optimizer = torch.optim.Adam(model.parameters())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import the Early Stopping Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import EarlyStopping\n",
    "from pytorchtools import EarlyStopping"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train the Model using Early Stopping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, batch_size, patience, n_epochs):\n",
    "    \n",
    "    # to track the training loss as the model trains\n",
    "    train_losses = []\n",
    "    # to track the validation loss as the model trains\n",
    "    valid_losses = []\n",
    "    # to track the average training loss per epoch as the model trains\n",
    "    avg_train_losses = []\n",
    "    # to track the average validation loss per epoch as the model trains\n",
    "    avg_valid_losses = [] \n",
    "    \n",
    "    # initialize the early_stopping object\n",
    "    early_stopping = EarlyStopping(patience=patience, verbose=True)\n",
    "    \n",
    "    for epoch in range(1, n_epochs + 1):\n",
    "\n",
    "        ###################\n",
    "        # train the model #\n",
    "        ###################\n",
    "        model.train() # prep model for training\n",
    "        for batch, (data, target) in enumerate(train_loader, 1):\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",
    "            # record training loss\n",
    "            train_losses.append(loss.item())\n",
    "\n",
    "        ######################    \n",
    "        # validate the model #\n",
    "        ######################\n",
    "        model.eval() # prep model for evaluation\n",
    "        for data, target in valid_loader:\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",
    "            # record validation loss\n",
    "            valid_losses.append(loss.item())\n",
    "\n",
    "        # print training/validation statistics \n",
    "        # calculate average loss over an epoch\n",
    "        train_loss = np.average(train_losses)\n",
    "        valid_loss = np.average(valid_losses)\n",
    "        avg_train_losses.append(train_loss)\n",
    "        avg_valid_losses.append(valid_loss)\n",
    "        \n",
    "        epoch_len = len(str(n_epochs))\n",
    "        \n",
    "        print_msg = (f'[{epoch:>{epoch_len}}/{n_epochs:>{epoch_len}}] ' +\n",
    "                     f'train_loss: {train_loss:.5f} ' +\n",
    "                     f'valid_loss: {valid_loss:.5f}')\n",
    "        \n",
    "        print(print_msg)\n",
    "        \n",
    "        # clear lists to track next epoch\n",
    "        train_losses = []\n",
    "        valid_losses = []\n",
    "        \n",
    "        # early_stopping needs the validation loss to check if it has decresed, \n",
    "        # and if it has, it will make a checkpoint of the current model\n",
    "        early_stopping(valid_loss, model)\n",
    "        \n",
    "        if early_stopping.early_stop:\n",
    "            print(\"Early stopping\")\n",
    "            break\n",
    "        \n",
    "    # load the last checkpoint with the best model\n",
    "    model.load_state_dict(torch.load('checkpoint.pt'))\n",
    "\n",
    "    return  model, avg_train_losses, avg_valid_losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[  1/100] train_loss: 0.84499 valid_loss: 0.29977\n",
      "[  2/100] train_loss: 0.36182 valid_loss: 0.22742\n",
      "Validation loss decreased (inf --> 0.227419).  Saving model ...\n",
      "[  3/100] train_loss: 0.29205 valid_loss: 0.19163\n",
      "Validation loss decreased (0.227419 --> 0.191628).  Saving model ...\n",
      "[  4/100] train_loss: 0.25390 valid_loss: 0.16771\n",
      "Validation loss decreased (0.191628 --> 0.167714).  Saving model ...\n",
      "[  5/100] train_loss: 0.22671 valid_loss: 0.15422\n",
      "Validation loss decreased (0.167714 --> 0.154222).  Saving model ...\n",
      "[  6/100] train_loss: 0.20862 valid_loss: 0.14546\n",
      "Validation loss decreased (0.154222 --> 0.145459).  Saving model ...\n",
      "[  7/100] train_loss: 0.19482 valid_loss: 0.13821\n",
      "Validation loss decreased (0.145459 --> 0.138206).  Saving model ...\n",
      "[  8/100] train_loss: 0.18431 valid_loss: 0.13398\n",
      "Validation loss decreased (0.138206 --> 0.133979).  Saving model ...\n",
      "[  9/100] train_loss: 0.17554 valid_loss: 0.12953\n",
      "Validation loss decreased (0.133979 --> 0.129535).  Saving model ...\n",
      "[ 10/100] train_loss: 0.16785 valid_loss: 0.12202\n",
      "Validation loss decreased (0.129535 --> 0.122023).  Saving model ...\n",
      "[ 11/100] train_loss: 0.16202 valid_loss: 0.12249\n",
      "EarlyStopping counter: 1 out of 20\n",
      "[ 12/100] train_loss: 0.15300 valid_loss: 0.11852\n",
      "Validation loss decreased (0.122023 --> 0.118516).  Saving model ...\n",
      "[ 13/100] train_loss: 0.14965 valid_loss: 0.11560\n",
      "Validation loss decreased (0.118516 --> 0.115598).  Saving model ...\n",
      "[ 14/100] train_loss: 0.14680 valid_loss: 0.11387\n",
      "Validation loss decreased (0.115598 --> 0.113867).  Saving model ...\n",
      "[ 15/100] train_loss: 0.13988 valid_loss: 0.11728\n",
      "EarlyStopping counter: 1 out of 20\n",
      "[ 16/100] train_loss: 0.13641 valid_loss: 0.11269\n",
      "Validation loss decreased (0.113867 --> 0.112686).  Saving model ...\n",
      "[ 17/100] train_loss: 0.12957 valid_loss: 0.11237\n",
      "Validation loss decreased (0.112686 --> 0.112374).  Saving model ...\n",
      "[ 18/100] train_loss: 0.12862 valid_loss: 0.11198\n",
      "Validation loss decreased (0.112374 --> 0.111975).  Saving model ...\n",
      "[ 19/100] train_loss: 0.12581 valid_loss: 0.10924\n",
      "Validation loss decreased (0.111975 --> 0.109242).  Saving model ...\n",
      "[ 20/100] train_loss: 0.12171 valid_loss: 0.10836\n",
      "Validation loss decreased (0.109242 --> 0.108363).  Saving model ...\n",
      "[ 21/100] train_loss: 0.12191 valid_loss: 0.10922\n",
      "EarlyStopping counter: 1 out of 20\n",
      "[ 22/100] train_loss: 0.11935 valid_loss: 0.10976\n",
      "EarlyStopping counter: 2 out of 20\n",
      "[ 23/100] train_loss: 0.11901 valid_loss: 0.11053\n",
      "EarlyStopping counter: 3 out of 20\n",
      "[ 24/100] train_loss: 0.11420 valid_loss: 0.10901\n",
      "EarlyStopping counter: 4 out of 20\n",
      "[ 25/100] train_loss: 0.11089 valid_loss: 0.10837\n",
      "EarlyStopping counter: 5 out of 20\n",
      "[ 26/100] train_loss: 0.11008 valid_loss: 0.10944\n",
      "EarlyStopping counter: 6 out of 20\n",
      "[ 27/100] train_loss: 0.10801 valid_loss: 0.10665\n",
      "Validation loss decreased (0.108363 --> 0.106647).  Saving model ...\n",
      "[ 28/100] train_loss: 0.10433 valid_loss: 0.10248\n",
      "Validation loss decreased (0.106647 --> 0.102475).  Saving model ...\n",
      "[ 29/100] train_loss: 0.10323 valid_loss: 0.10621\n",
      "EarlyStopping counter: 1 out of 20\n",
      "[ 30/100] train_loss: 0.10484 valid_loss: 0.10775\n",
      "EarlyStopping counter: 2 out of 20\n",
      "[ 31/100] train_loss: 0.09985 valid_loss: 0.10616\n",
      "EarlyStopping counter: 3 out of 20\n",
      "[ 32/100] train_loss: 0.09898 valid_loss: 0.10479\n",
      "EarlyStopping counter: 4 out of 20\n",
      "[ 33/100] train_loss: 0.10062 valid_loss: 0.10576\n",
      "EarlyStopping counter: 5 out of 20\n",
      "[ 34/100] train_loss: 0.09704 valid_loss: 0.10770\n",
      "EarlyStopping counter: 6 out of 20\n",
      "[ 35/100] train_loss: 0.09850 valid_loss: 0.10542\n",
      "EarlyStopping counter: 7 out of 20\n",
      "[ 36/100] train_loss: 0.09561 valid_loss: 0.10619\n",
      "EarlyStopping counter: 8 out of 20\n",
      "[ 37/100] train_loss: 0.09381 valid_loss: 0.10745\n",
      "EarlyStopping counter: 9 out of 20\n",
      "[ 38/100] train_loss: 0.09363 valid_loss: 0.10487\n",
      "EarlyStopping counter: 10 out of 20\n",
      "[ 39/100] train_loss: 0.09263 valid_loss: 0.10763\n",
      "EarlyStopping counter: 11 out of 20\n",
      "[ 40/100] train_loss: 0.09234 valid_loss: 0.10778\n",
      "EarlyStopping counter: 12 out of 20\n",
      "[ 41/100] train_loss: 0.08485 valid_loss: 0.10319\n",
      "EarlyStopping counter: 13 out of 20\n",
      "[ 42/100] train_loss: 0.09105 valid_loss: 0.10305\n",
      "EarlyStopping counter: 14 out of 20\n",
      "[ 43/100] train_loss: 0.08963 valid_loss: 0.10952\n",
      "EarlyStopping counter: 15 out of 20\n",
      "[ 44/100] train_loss: 0.08887 valid_loss: 0.10615\n",
      "EarlyStopping counter: 16 out of 20\n",
      "[ 45/100] train_loss: 0.08704 valid_loss: 0.10870\n",
      "EarlyStopping counter: 17 out of 20\n",
      "[ 46/100] train_loss: 0.08477 valid_loss: 0.10877\n",
      "EarlyStopping counter: 18 out of 20\n",
      "[ 47/100] train_loss: 0.08397 valid_loss: 0.10682\n",
      "EarlyStopping counter: 19 out of 20\n",
      "[ 48/100] train_loss: 0.08630 valid_loss: 0.10565\n",
      "EarlyStopping counter: 20 out of 20\n",
      "Early stopping\n"
     ]
    }
   ],
   "source": [
    "batch_size = 256\n",
    "n_epochs = 100\n",
    "\n",
    "train_loader, test_loader, valid_loader = create_datasets(batch_size)\n",
    "\n",
    "# early stopping patience; how long to wait after last time validation loss improved.\n",
    "patience = 20\n",
    "\n",
    "model, train_loss, valid_loss = train_model(model, batch_size, patience, n_epochs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualizing the Loss and the Early Stopping Checkpoint\n",
    "From the plot we can see that the last Early Stopping Checkpoint was saved right before the model started to overfit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 38132 characters]\n",
      "text/plain": [
       "<Figure size 720x576 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# visualize the loss as the network trained\n",
    "fig = plt.figure(figsize=(10,8))\n",
    "plt.plot(range(1,len(train_loss)+1),train_loss, label='Training Loss')\n",
    "plt.plot(range(1,len(valid_loss)+1),valid_loss,label='Validation Loss')\n",
    "\n",
    "# find position of lowest validation loss\n",
    "minposs = valid_loss.index(min(valid_loss))+1 \n",
    "plt.axvline(minposs, linestyle='--', color='r',label='Early Stopping Checkpoint')\n",
    "\n",
    "plt.xlabel('epochs')\n",
    "plt.ylabel('loss')\n",
    "plt.ylim(0, 0.5) # consistent scale\n",
    "plt.xlim(0, len(train_loss)+1) # consistent scale\n",
    "plt.grid(True)\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "fig.savefig('loss_plot.png', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test the Trained Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.085796\n",
      "\n",
      "Test Accuracy of     0: 98% (968/979)\n",
      "Test Accuracy of     1: 98% (1121/1133)\n",
      "Test Accuracy of     2: 97% (1008/1030)\n",
      "Test Accuracy of     3: 97% (986/1008)\n",
      "Test Accuracy of     4: 97% (955/980)\n",
      "Test Accuracy of     5: 97% (867/890)\n",
      "Test Accuracy of     6: 97% (932/956)\n",
      "Test Accuracy of     7: 97% (997/1027)\n",
      "Test Accuracy of     8: 96% (939/973)\n",
      "Test Accuracy of     9: 95% (967/1008)\n",
      "\n",
      "Test Accuracy (Overall): 97% (9740/9984)\n"
     ]
    }
   ],
   "source": [
    "# initialize lists to monitor test loss and accuracy\n",
    "test_loss = 0.0\n",
    "class_correct = list(0. for i in range(10))\n",
    "class_total = list(0. for i in range(10))\n",
    "\n",
    "model.eval() # prep model for evaluation\n",
    "\n",
    "for data, target in test_loader:\n",
    "    if len(target.data) != batch_size:\n",
    "        break\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 test loss \n",
    "    test_loss += loss.item()*data.size(0)\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",
    "# calculate and print avg test loss\n",
    "test_loss = test_loss/len(test_loader.dataset)\n",
    "print('Test Loss: {:.6f}\\n'.format(test_loss))\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": [
    "### Visualize Sample Test Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 27400 characters]\n",
      "text/plain": [
       "<Figure size 1800x288 with 20 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# obtain one batch of test images\n",
    "dataiter = iter(test_loader)\n",
    "images, labels = dataiter.next()\n",
    "\n",
    "# get sample outputs\n",
    "output = model(images)\n",
    "# convert output probabilities to predicted class\n",
    "_, preds = torch.max(output, 1)\n",
    "# prep images for display\n",
    "images = images.numpy()\n",
    "\n",
    "# plot the images in the batch, along with predicted and true labels\n",
    "fig = plt.figure(figsize=(25, 4))\n",
    "for idx in np.arange(20):\n",
    "    ax = fig.add_subplot(2, 20/2, idx+1, xticks=[], yticks=[])\n",
    "    ax.imshow(np.squeeze(images[idx]), cmap='gray')\n",
    "    ax.set_title(\"{} ({})\".format(str(preds[idx].item()), str(labels[idx].item())),\n",
    "                 color=(\"green\" if preds[idx]==labels[idx] else \"red\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The MNIST training example code is mainly taken from the [mnist_mlp_solution_with_validation.ipynb](https://github.com/udacity/deep-learning-v2-pytorch/blob/master/convolutional-neural-networks/mnist-mlp/mnist_mlp_solution_with_validation.ipynb) notebook from the [deep-learning-v2-pytorch repository](https://github.com/udacity/deep-learning-v2-pytorch) made by [Udacity](https://www.udacity.com/), and has been fitted with my early stopping code. "
   ]
  }
 ],
 "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.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
