{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# **Image Classifier with PyTorch**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 0. Set Up\n",
    "\n",
    "For this tutorial, we will be using ``numpy`` to generate data and all operations, ``matplotlib`` to render and plot images and of course ``nets`` for ``Tensor`` operations and computational graph creation. We will set the ``SEED`` for random numbers from ``numpy``, so each time we reload the notebook the results will be the same as numbers will be generated in the same way. Finally, we set an autoreload with ``%load_ext autoreload`` and ``%autoreload 2`` for debug mode purposes. This enable the notebook to reload the modules each time we run a cell, so it keeps everything up-to-date. More info [on StackOverflow](http://stackoverflow.com/questions/1907993/autoreload-of-modules-in-ipython).\n",
    "\n",
    "Let's load the packages we need to run our tests:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random as rd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# This is a bit of magic to make matplotlib figures appear inline in the notebook\n",
    "# rather than in a new window.\n",
    "%matplotlib inline\n",
    "plt.rcParams['figure.figsize'] = (10.0, 8.0) # set default size of plots\n",
    "plt.rcParams['image.interpolation'] = 'nearest'\n",
    "plt.style.use('seaborn-darkgrid')\n",
    "# Change the font ?\n",
    "font = {'size'   : 15}\n",
    "plt.rc('font', **font)\n",
    "\n",
    "# Some more magic so that the notebook will reload external python modules;\n",
    "# see http://stackoverflow.com/questions/1907993/autoreload-of-modules-in-ipython\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's fix the seed\n",
    "\n",
    "SEED = 42\n",
    "rd.seed(SEED)\n",
    "np.random.seed(SEED)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Import the Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to .data\\MNIST\\raw\\train-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "25e51451b563449abba66d6ff8e8b516",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', max=1.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting .data\\MNIST\\raw\\train-images-idx3-ubyte.gz to .data\\MNIST\\raw\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to .data\\MNIST\\raw\\train-labels-idx1-ubyte.gz\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "969bee51206f43638cd4662306efc99e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', max=1.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting .data\\MNIST\\raw\\train-labels-idx1-ubyte.gz to .data\\MNIST\\raw\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to .data\\MNIST\\raw\\t10k-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6d383dde34514907a6a6062372e46c30",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', max=1.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting .data\\MNIST\\raw\\t10k-images-idx3-ubyte.gz to .data\\MNIST\\raw\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to .data\\MNIST\\raw\\t10k-labels-idx1-ubyte.gz\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "83ef2266e89845448d42347fa0a65524",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', max=1.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting .data\\MNIST\\raw\\t10k-labels-idx1-ubyte.gz to .data\\MNIST\\raw\n",
      "Processing...\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import random as rd\n",
    "import torch\n",
    "import torchvision\n",
    "import matplotlib.pyplot as plt\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "# Normalize the data !\n",
    "transform = transforms.Compose([transforms.ToTensor(),\n",
    "                              transforms.Normalize((0.5,), (0.5,)),\n",
    "                              ])\n",
    "\n",
    "# Download the dataset on the default 'root' folder, '.data'\n",
    "root = '.data'\n",
    "train_data = datasets.MNIST(root, download=True, train=True, transform=transform)\n",
    "test_data = datasets.MNIST(root, download=True, train=False, transform=transform)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Load the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sanity Check\n",
      "torch.Size([60000, 28, 28])\n",
      "torch.Size([60000])\n"
     ]
    }
   ],
   "source": [
    "print('Sanity Check')\n",
    "print(train_data.data.shape)\n",
    "print(train_data.targets.shape)\n",
    "\n",
    "BATCH_SIZE = 64\n",
    "train_loader = torch.utils.data.DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 38128 characters]\n",
      "text/plain": [
       "<Figure size 1440x720 with 32 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Configure the mosaic to display images\n",
    "fig = plt.figure(figsize = (20, 10))\n",
    "columns = 8\n",
    "rows = 4\n",
    "\n",
    "# Create the mosaic\n",
    "for i in range(1, columns * rows + 1):\n",
    "    \n",
    "    # Choose randomly an image\n",
    "    image_index = rd.randint(0, len(train_data))\n",
    "    image, label = train_data.data[image_index], train_data.targets[image_index]\n",
    "    \n",
    "    # Plot one image in the grid\n",
    "    fig.add_subplot(rows, columns, i)\n",
    "    plt.imshow(image.numpy().reshape(28, 28))\n",
    "    \n",
    "    #Remove axis\n",
    "    plt.axis('off')\n",
    "    \n",
    "# Render\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Build a model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Classifier(\n",
       "  (layer1): Linear(in_features=784, out_features=100, bias=True)\n",
       "  (layer2): Linear(in_features=100, out_features=100, bias=True)\n",
       "  (layer3): Linear(in_features=100, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "# Just a simple model\n",
    "class Classifier(nn.Module):\n",
    "\n",
    "    def __init__(self, input_dim, hidden_dim, output_dim):\n",
    "        super().__init__()\n",
    "        self.layer1 = nn.Linear(input_dim, hidden_dim)\n",
    "        self.layer2 = nn.Linear(hidden_dim, hidden_dim)\n",
    "        self.layer3 = nn.Linear(hidden_dim, output_dim)\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        out1 = nn.functional.relu(self.layer1(inputs))\n",
    "        out2 = nn.functional.relu(self.layer2(out1))\n",
    "        return self.layer3(out2)\n",
    "    \n",
    "model = Classifier(784, 100, 10)\n",
    "\n",
    "# Let's check the architecture\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Criterion & Optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.optim import SGD\n",
    "from torch.nn import CrossEntropyLoss\n",
    "from nets.nn.utils import one_hot\n",
    "\n",
    "# How much do we update the parameters\n",
    "LEARNING_RATE = 0.01\n",
    "\n",
    "optimizer = SGD(model.parameters(), lr=LEARNING_RATE)\n",
    "criterion = CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5. Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 0/50 | loss: 2.950E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 1/50 | loss: 2.731E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 2/50 | loss: 2.539E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 3/50 | loss: 2.355E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 4/50 | loss: 2.186E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 5/50 | loss: 2.039E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 6/50 | loss: 1.906E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 7/50 | loss: 1.781E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 8/50 | loss: 1.671E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 9/50 | loss: 1.570E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 10/50 | loss: 1.478E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 11/50 | loss: 1.396E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 12/50 | loss: 1.321E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 13/50 | loss: 1.250E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 14/50 | loss: 1.192E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 15/50 | loss: 1.134E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 16/50 | loss: 1.080E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 17/50 | loss: 1.030E-01\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 18/50 | loss: 9.879E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 19/50 | loss: 9.442E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 20/50 | loss: 9.095E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 21/50 | loss: 8.704E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 22/50 | loss: 8.378E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 23/50 | loss: 8.084E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 24/50 | loss: 7.758E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 25/50 | loss: 7.489E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 26/50 | loss: 7.195E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 27/50 | loss: 6.931E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 28/50 | loss: 6.699E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 29/50 | loss: 6.439E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 30/50 | loss: 6.271E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 31/50 | loss: 6.081E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 32/50 | loss: 5.878E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 33/50 | loss: 5.635E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 34/50 | loss: 5.474E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 35/50 | loss: 5.296E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 36/50 | loss: 5.073E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 37/50 | loss: 4.975E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 38/50 | loss: 4.800E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 39/50 | loss: 4.658E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 40/50 | loss: 4.525E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 41/50 | loss: 4.391E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 42/50 | loss: 4.217E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 43/50 | loss: 4.089E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 44/50 | loss: 3.974E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 45/50 | loss: 3.851E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 46/50 | loss: 3.751E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 47/50 | loss: 3.594E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 48/50 | loss: 3.543E-02\n",
      "Training: 100% | [==================================================] | Done !\n",
      "epoch: 49/50 | loss: 3.388E-02\n",
      "\n",
      "time: 1165.0111677646637\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "from nets.utils import progress_bar\n",
    "\n",
    "# Training iterations\n",
    "EPOCHS = 50\n",
    "\n",
    "#Display the time during training\n",
    "start = time.time()\n",
    "# Record the loss\n",
    "history = []\n",
    "\n",
    "# Run the simulation EPOCHS times\n",
    "for epoch in range(EPOCHS):\n",
    "    \n",
    "    # Keep track of the loss at each epoch\n",
    "    epoch_loss = 0.0\n",
    "    epoch_start = time.time()\n",
    "    \n",
    "    # Number of batches\n",
    "    N  = len(train_loader)\n",
    "\n",
    "    # Predict for all batches\n",
    "    for idx, batch in enumerate(train_loader):\n",
    "        \n",
    "        # Gradients accumulates, therefore we need to set them to zero at each iteration\n",
    "        model.zero_grad()\n",
    "        \n",
    "        # Get the images and labels from one batch\n",
    "        image, label = batch\n",
    "        # Let's not forget to reshape the images from (BATCH_SIE, 1, 28, 28) to (BATCH_SIZE, 1 * 28 * 28)\n",
    "        image = image.reshape(-1, 784)\n",
    "        \n",
    "        # Run the forward pass by calling the model on our images\n",
    "        predictions = model(image.float())\n",
    "        # Get our loss ie how far our predictions are from reality\n",
    "        loss = criterion(predictions, label.long())\n",
    "        \n",
    "        # Compute the gradient\n",
    "        loss.backward()\n",
    "        # Update the parameters\n",
    "        optimizer.step()\n",
    "        \n",
    "        # Record the loss for plotting\n",
    "        epoch_loss += loss.item()\n",
    "        # Let's plot a progress bar in the console\n",
    "        progress_bar(idx, N-1, \"Training:\")\n",
    "    \n",
    "    # Update the history of all previous loss\n",
    "    history.append(epoch_loss / N)\n",
    "    \n",
    "    # Check the mean loss for this loop\n",
    "    print(f\"epoch: {epoch}/{EPOCHS} | loss: {epoch_loss / N:1.3E}\")\n",
    "\n",
    "# Duration of the total training loop\n",
    "print('\\ntime:', time.time() - start)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x21e7bd0b988>]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 12316 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(history)"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
