labs/08_frameworks/Fashion_MNIST_classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fashion MNIST classification using PyTorch\n",
    "\n",
    "In this notebook we will try to classify the Fashion MNIST dataset\n",
    "(https://github.com/zalandoresearch/fashion-mnist) using VGG-like architectures (https://arxiv.org/abs/1409.1556). This notebook is inspired from the MNIST example from PyTorch (https://github.com/pytorch/examples/tree/master/mnist), and introduce tricks to automatically tune and schedule the learning rate for SGD (see this course's slides, https://arxiv.org/abs/1506.01186, and FastAI course for example http://fastai.org).\n",
    "\n",
    "## Fashion MNIST\n",
    "\n",
    "This 10 class dataset is a drop-in replacement for MNIST with clothes instead of digits. MNI is arguably overused in the ML community nowadays. It is subtancially harder to classify.\n",
    "\n",
    "![fashion_mnist](fashion-mnist-sprite.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's import a few functions first"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from matplotlib.cm import get_cmap\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "from torchvision import datasets\n",
    "from torchvision.transforms import transforms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some system/model hyper parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cuda = False\n",
    "batch_size = 128\n",
    "\n",
    "kwargs = {'num_workers': 1, 'pin_memory': True} if cuda else {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "english_labels = [\"T-shirt/top\",\n",
    "                  \"Trouser\",\n",
    "                  \"Pullover\",\n",
    "                  \"Dress\",\n",
    "                  \"Coat\",\n",
    "                  \"Sandal\",\n",
    "                  \"Shirt\",\n",
    "                  \"Sneaker\",\n",
    "                  \"Bag\",\n",
    "                  \"Ankle boot\"]\n",
    "\n",
    "train_data = datasets.FashionMNIST('data', train=True, download=True,\n",
    "                                   transform=transforms.Compose([\n",
    "                                       transforms.ToTensor(),\n",
    "                                   ]))\n",
    "train_loader = DataLoader(train_data, batch_size=128, shuffle=False, **kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets's compute the average mean and std of the train images. We will\n",
    "use them for normalizing data later on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_samples_seen = 0.\n",
    "mean = 0\n",
    "std = 0\n",
    "for train_batch, train_target in train_loader:\n",
    "    batch_size = train_batch.shape[0]\n",
    "    train_batch = train_batch.view(batch_size, -1)\n",
    "    this_mean = torch.mean(train_batch, dim=1)\n",
    "    this_std = torch.sqrt(\n",
    "        torch.mean((train_batch - this_mean[:, None]) ** 2, dim=1))\n",
    "    mean += torch.sum(this_mean, dim=0)\n",
    "    std += torch.sum(this_std, dim=0)\n",
    "    n_samples_seen += batch_size\n",
    "\n",
    "mean /= n_samples_seen\n",
    "std /= n_samples_seen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(mean, std)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now reload the data with a further `Normalize` transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = datasets.FashionMNIST('data', train=True, download=False,\n",
    "                                   transform=transforms.Compose([\n",
    "                                       transforms.ToTensor(),\n",
    "                                       transforms.Normalize(mean=mean.view(1),\n",
    "                                                            std=std.view(1))]))\n",
    "\n",
    "test_data = datasets.FashionMNIST('data', train=False, download=True,\n",
    "                                  transform=transforms.Compose([\n",
    "                                      transforms.ToTensor(),\n",
    "                                      transforms.Normalize(mean=mean.view(1),\n",
    "                                                           std=std.view(1))]))\n",
    "\n",
    "train_loader = DataLoader(train_data, batch_size=32, shuffle=True, **kwargs)\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(test_data, batch_size=32,\n",
    "                                          shuffle=False, **kwargs)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define a very simple model, suitable for CPU training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Model, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 10, kernel_size=(3, 3), padding=1)\n",
    "        self.conv2 = nn.Conv2d(10, 20, kernel_size=(3, 3), padding=1)\n",
    "        self.dropout_2d = nn.Dropout2d(p=0.25)\n",
    "        self.fc1 = nn.Linear(7 * 7 * 20, 128)\n",
    "        self.dropout = nn.Dropout(p=0.5)\n",
    "        self.fc2 = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.dropout_2d(F.max_pool2d(self.conv1(x), kernel_size=2))\n",
    "        x = self.dropout_2d(F.max_pool2d(self.conv2(x), kernel_size=2))\n",
    "        x = x.view(-1, 7 * 7 * 20)  # flatten / reshape\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = self.dropout(x)\n",
    "        x = self.fc2(x)\n",
    "        return F.log_softmax(x, dim=1)\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        self.conv1.reset_parameters()\n",
    "        self.conv2.reset_parameters()\n",
    "        self.fc1.reset_parameters()\n",
    "        self.fc2.reset_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercices\n",
    "\n",
    "- Define a VGG-like model: add more convolutional and max pooling layers to increase the number of channels progressively while decreasing the dimensions of the feature maps with max pooling.\n",
    "- Try to use Adam instead of SGD in conjunction with the `find_lr` heuristic and the cosine learning rate schedule below;\n",
    "- (optional) Try data augmentation (horizontal flips, random crops, cutout...);\n",
    "- (optional) Implement the [mixup stochastic label interpolation](https://arxiv.org/abs/1710.09412);\n",
    "- (optional) Try to use batch-normalization;\n",
    "- (optional) Implement skip-connections.\n",
    "\n",
    "See how you compare to other approaches:\n",
    "- https://github.com/zalandoresearch/fashion-mnist\n",
    "- https://www.kaggle.com/zalando-research/fashionmnist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/vgg.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's test our model on CPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model()\n",
    "img, target = train_data[0]\n",
    "# n_channel, width, height\n",
    "print(img.shape)\n",
    "\n",
    "# First dimension should contain batch_size\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can plot a training image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "ax.imshow(img[0].numpy(), cmap=get_cmap('gray'))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first dimension of the input data should contain the batch size (due to `torch.nn` API)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = model(img[None, :])\n",
    "print(target, english_labels[target])\n",
    "print(pred)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define a train loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, optimizer, train_loader, epoch):\n",
    "    model.train()\n",
    "    train_loss = 0\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        if cuda:\n",
    "            data, target = data.cuda(), target.cuda()\n",
    "        batch_size = data.shape[0]\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        loss = F.nll_loss(output, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        train_loss += loss.item() * batch_size\n",
    "\n",
    "        if batch_idx % log_interval == 0:\n",
    "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "                epoch + 1, batch_idx * len(data), len(train_loader.dataset),\n",
    "                100. * batch_idx / len(train_loader), loss.item()))\n",
    "    train_loss /= len(test_loader.dataset)\n",
    "    return train_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And a test function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(model, test_loader):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            if cuda:\n",
    "                data, target = data.cuda(), target.cuda()\n",
    "            output = model(data)\n",
    "            test_loss += F.nll_loss(output, target, size_average=False).item()\n",
    "            # sum up batch loss\n",
    "            _, pred = output.data.max(dim=1)\n",
    "            # get the index of the max log-probability\n",
    "            correct += torch.sum(pred == target.data.long()).item()\n",
    "\n",
    "        test_loss /= len(test_loader.dataset)\n",
    "        test_accuracy = float(correct) / len(test_loader.dataset)\n",
    "        print('\\nTest set: Average loss: {:.4f},'\n",
    "              ' Accuracy: {}/{} ({:.0f}%)\\n'.format(\n",
    "            test_loss, correct, len(test_loader.dataset),\n",
    "            100. * test_accuracy))\n",
    "    return test_loss, test_accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `find_lr` function provides a learning rate for SGD or Adam, following heuristics from https://arxiv.org/abs/1506.01186:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def loop_loader(data_loader):\n",
    "    while True:\n",
    "        for elem in data_loader:\n",
    "            yield elem\n",
    "\n",
    "def find_lr(model, train_loader, init_lr, max_lr, steps, n_batch_per_step=30):\n",
    "    optimizer = torch.optim.SGD(model.parameters(), lr=init_lr)\n",
    "    current_lr = init_lr\n",
    "    best_lr = current_lr\n",
    "    best_loss = float('inf')\n",
    "    lr_step = (max_lr - init_lr) / steps\n",
    "\n",
    "    loader = loop_loader(train_loader)\n",
    "    for i in range(steps):\n",
    "        mean_loss = 0\n",
    "        n_seen_samples = 0\n",
    "        for j, (data, target) in enumerate(loader):\n",
    "            if j > n_batch_per_step:\n",
    "                break\n",
    "            optimizer.zero_grad()\n",
    "            if cuda:\n",
    "                data = data.cuda()\n",
    "                target = target.cuda()\n",
    "            output = model(data)\n",
    "            loss = F.nll_loss(output, target)\n",
    "            loss.backward()\n",
    "            mean_loss += loss.item() * data.shape[0]\n",
    "            n_seen_samples += data.shape[0]\n",
    "            optimizer.step()\n",
    "\n",
    "        mean_loss /= n_seen_samples\n",
    "        print('Step %i, current LR: %f, loss %f' % (i, current_lr, mean_loss))\n",
    "            \n",
    "        if np.isnan(mean_loss) or mean_loss > best_loss * 4:\n",
    "            return best_lr / 4\n",
    "        \n",
    "        if mean_loss < best_loss:\n",
    "            best_loss = mean_loss\n",
    "            best_lr = current_lr\n",
    "\n",
    "        current_lr += lr_step\n",
    "        optimizer.param_groups[0]['lr'] = current_lr\n",
    "\n",
    "    return best_lr / 4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load our model on the GPU if required. We then define an optimizer and a learning rate scheduler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_interval = 100\n",
    "epochs = 12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if cuda:\n",
    "    model.cuda()\n",
    "\n",
    "model.reset_parameters()\n",
    "lr = find_lr(model, train_loader, 1e-4, 1, 100, 30)\n",
    "model.reset_parameters()\n",
    "\n",
    "print('Best LR', lr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's train our model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.SGD(model.parameters(), lr=lr)\n",
    "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,\n",
    "                                                       T_max=3,\n",
    "                                                       last_epoch=-1)\n",
    "\n",
    "logs = {'epoch': [], 'train_loss': [], 'test_loss': [],\n",
    "        'test_accuracy': [], 'lr': []}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for epoch in range(epochs):\n",
    "    train_loss = train(model, optimizer, train_loader, epoch)\n",
    "    test_loss, test_accuracy = test(model, test_loader)\n",
    "    logs['epoch'].append(epoch)\n",
    "    logs['train_loss'].append(train_loss)\n",
    "    logs['test_loss'].append(test_loss)\n",
    "    logs['test_accuracy'].append(test_accuracy)\n",
    "    logs['lr'].append(optimizer.param_groups[0]['lr'])\n",
    "    scheduler.step(epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}