mnist-complete.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MNIST classifier with PyTorch\n",
    "\n",
    "The MNIST dataset (http://yann.lecun.com/exdb/mnist/) contains 60k labeled images of handwritten digits.\n",
    "Using the PyTorch framework, we'll write a classifier to classify the given images into their digit class.\n",
    "\n",
    "We'll proceed as follows:\n",
    "- Create the data pipeline\n",
    "- Create model\n",
    "- Train model\n",
    "- Test model\n",
    "- Put everything together"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dataset is located in `./data` and is contained in 4 binary files `train-images`, `train-labels`, `test-images`, `test-labels`. As the data is in binary format it has to be read and decoded. We'll use the following `read_image_file` and `read_label_file` helper functions to do so easily:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import codecs\n",
    "from torch.utils.data import Dataset\n",
    "from torchvision import transforms\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "########## HELPER FUNCTIONS ##########\n",
    "# Source: [https://pytorch.org/docs/stable/_modules/torchvision/datasets/mnist.html]\n",
    "######################################\n",
    "\n",
    "def read_label_file(path):\n",
    "    with open(path, 'rb') as f:\n",
    "        x = read_sn3_pascalvincent_tensor(f, strict=False)\n",
    "    assert(x.dtype == torch.uint8)\n",
    "    assert(x.ndimension() == 1)\n",
    "    return x.long()\n",
    "\n",
    "\n",
    "def read_image_file(path):\n",
    "    with open(path, 'rb') as f:\n",
    "        x = read_sn3_pascalvincent_tensor(f, strict=False)\n",
    "    assert(x.dtype == torch.uint8)\n",
    "    assert(x.ndimension() == 3)\n",
    "    return x\n",
    "\n",
    "def get_int(b):\n",
    "  return int(codecs.encode(b, 'hex'), 16)\n",
    "\n",
    "def open_maybe_compressed_file(path):\n",
    "  \"\"\"Return a file object that possibly decompresses 'path' on the fly.\n",
    "      Decompression occurs when argument `path` is a string and ends with '.gz' or '.xz'.\n",
    "  \"\"\"\n",
    "  if not isinstance(path, torch._six.string_classes):\n",
    "      return path\n",
    "  if path.endswith('.gz'):\n",
    "      import gzip\n",
    "      return gzip.open(path, 'rb')\n",
    "  if path.endswith('.xz'):\n",
    "      import lzma\n",
    "      return lzma.open(path, 'rb')\n",
    "  return open(path, 'rb')\n",
    "\n",
    "def read_sn3_pascalvincent_tensor(path, strict=True):\n",
    "    \"\"\"Read a SN3 file in \"Pascal Vincent\" format (Lush file 'libidx/idx-io.lsh').\n",
    "       Argument may be a filename, compressed filename, or file object.\n",
    "    \"\"\"\n",
    "    # typemap\n",
    "    if not hasattr(read_sn3_pascalvincent_tensor, 'typemap'):\n",
    "        read_sn3_pascalvincent_tensor.typemap = {\n",
    "            8: (torch.uint8, np.uint8, np.uint8),\n",
    "            9: (torch.int8, np.int8, np.int8),\n",
    "            11: (torch.int16, np.dtype('>i2'), 'i2'),\n",
    "            12: (torch.int32, np.dtype('>i4'), 'i4'),\n",
    "            13: (torch.float32, np.dtype('>f4'), 'f4'),\n",
    "            14: (torch.float64, np.dtype('>f8'), 'f8')}\n",
    "    # read\n",
    "    with open_maybe_compressed_file(path) as f:\n",
    "        data = f.read()\n",
    "    # parse\n",
    "    magic = get_int(data[0:4])\n",
    "    nd = magic % 256\n",
    "    ty = magic // 256\n",
    "    assert nd >= 1 and nd <= 3\n",
    "    assert ty >= 8 and ty <= 14\n",
    "    m = read_sn3_pascalvincent_tensor.typemap[ty]\n",
    "    s = [get_int(data[4 * (i + 1): 4 * (i + 2)]) for i in range(nd)]\n",
    "    parsed = np.frombuffer(data, dtype=m[1], offset=(4 * (nd + 1)))\n",
    "    assert parsed.shape[0] == np.prod(s) or not strict\n",
    "    return torch.from_numpy(parsed.astype(m[2], copy=False)).view(*s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Creating the data pipeline\n",
    "We first start by creating a data pipeline for our dataset. This will allow us to efficiently feed data to our model as it trains. To do so, we inherit from the `Dataset` class and override the `__len__(self)` and `__getitem(self, index)__` methods. We initialise the class with \n",
    "- the path to the data\n",
    "- whether to load the train or test data\n",
    "- a list of transforms for input (optional)\n",
    "- a list of transforms for labels (optional)\n",
    "\n",
    "The `__init__` method should load the right dataset and the `__getitem__` method should apply transforms if necessary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MNISTDataset(Dataset):\n",
    "  def __init__(self, path, test=False, transform=None, label_transform=None):\n",
    "    if test:\n",
    "      self.x = read_image_file(path + '/test-images')\n",
    "      self.y = read_label_file(path + '/test-labels')\n",
    "    else:\n",
    "      self.x = read_image_file(path + '/train-images')\n",
    "      self.y = read_label_file(path + '/train-labels')\n",
    "\n",
    "    assert(len(self.x) == len(self.y))\n",
    "\n",
    "    self.transform = transform\n",
    "    self.label_transform = label_transform\n",
    "\n",
    "  def __len__(self):\n",
    "    return len(self.x)\n",
    "\n",
    "  def __getitem__(self, idx):\n",
    "    sample = self.x[idx]\n",
    "    label = self.y[idx]\n",
    "\n",
    "    sample = Image.fromarray(sample.numpy(), mode='L')\n",
    "\n",
    "    if self.transform is not None:\n",
    "      sample = self.transform(sample)\n",
    "\n",
    "    if self.label_transform is not None:\n",
    "      label = self.label_transform(label)\n",
    "\n",
    "    return sample, label\n",
    "\n",
    "dataset = MNISTDataset('./data')\n",
    "print(dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our dataset class will now read the files and output pairs of PIL (Python Image Library) images and the corresponding label class (i.e. the image is a 4)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now create our dataloaders for both the train and test set. Note that since our Dataset outputs PIL images, those have to be converted to a PyTorch tensor using the `transforms.ToTensor()` transform. We define the `BATCH_SIZE` hyperparameter to define how big our batches are for gradient descent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 32\n",
    "DATA = './data'\n",
    "\n",
    "train_data = MNISTDataset(DATA, test=False, transform=transforms.ToTensor())\n",
    "trainloader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)\n",
    "\n",
    "test_data = MNISTDataset(DATA, test=True, transform=transforms.ToTensor())\n",
    "testloader = DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Creating the model\n",
    "\n",
    "Now that our data pipeline is up and running we can create our model. We take as input a 28x28 image which we flatten out to a 28x28=784 vector. As such our input layer will be 784 wide, we'll then use 1 64 hidden layer and a final 10 output layer. The output layer is 10 wide since we output one hot encoded vectors of class. Having a single output neuron ranging from 0 to 9 would be very difficult to train. In our case, if the second neuron is lit up, we classify it as a 1 (we start from 0).\n",
    "\n",
    "We will use the ReLU activation function for all layers and softmax for the output layer in order to get a probability distribution.\n",
    "\n",
    "Our model can now be created by inheriting from the nn.Module class and defining the `forward(self, x)` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class Model(nn.Module):\n",
    "  def __init__(self):\n",
    "    super(Model, self).__init__()\n",
    "\n",
    "    self.fc1 = nn.Linear(28 * 28, 64)\n",
    "    self.fc2 = nn.Linear(64, 64)\n",
    "    self.fc3 = nn.Linear(64, 10)\n",
    "\n",
    "  def forward(self, x):\n",
    "    x = F.relu(self.fc1(x))\n",
    "    x = F.relu(self.fc2(x))\n",
    "    x = F.softmax(self.fc3(x), dim=1)\n",
    "\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a tiny model and far from the most optimal one. It is however quite simple and does the job. Feel free to play around with different designs!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Train model\n",
    "\n",
    "Now that we have our data pipeline and our model we have to train it. We'll define the `train` function taking the following arguments:\n",
    "- epoch: current epoch number\n",
    "- model: model we're training\n",
    "- dataloader: dataloader to get the data from\n",
    "- optimizer: optimizer to run on the model\n",
    "- loss_fn: loss function to compute loss\n",
    "\n",
    "Using the dataloader, we iterate over batches, feed them in, compute the loss and backpropagate the gradients.\n",
    "Note that for every batch `optimizer.zero_grad()` should be called to reset the gradients otherwise they would accumulate indefinitely.\n",
    "\n",
    "The input data will be of the form `[batch_size, n_channels, height, width]`, so we'll get a tensor along the lines of `[32, 1, 28, 28]` (only 1 channel since the image is grayscale). However, our model expects a 728 flat input so we have to reshape the data into a tensor of dimension `[batch_size, 28 * 28]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LOG_INTERVAL=100\n",
    "\n",
    "def train(epoch, model, dataloader, optimizer, loss_fn):\n",
    "  for i, batch in enumerate(dataloader):\n",
    "    (inputs, labels) = batch\n",
    "\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    inputs = inputs.reshape(-1, 28 * 28)\n",
    "    outputs = model(inputs)\n",
    "\n",
    "    loss = loss_fn(outputs, labels)\n",
    "    loss.backward()\n",
    "\n",
    "    optimizer.step()\n",
    "\n",
    "    if i % LOG_INTERVAL == 0:\n",
    "      print('Epoch {}: [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "          epoch, i * len(inputs), len(dataloader.dataset),\n",
    "          100. * i / len(dataloader), loss.item()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Test model\n",
    "The dataset gives 10k images to test our model on. Much like the `train` function, we'll create a `test` function to test our model on unseen data. We can run the code under the `with torch.no_grad():` directive to tell PyTorch we won't be using gradients so it shouldn't waste time computing them.\n",
    "We wish to track the average loss and the percentage of correct predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(model, dataloader, loss_fn):\n",
    "  loss = 0\n",
    "  correct = 0\n",
    "  with torch.no_grad():\n",
    "    for x, y in dataloader:\n",
    "      x = x.reshape(-1, 28 * 28)\n",
    "      y_pred = model(x)\n",
    "\n",
    "      loss += loss_fn(y_pred, y).item()\n",
    "\n",
    "      # Transform from vector of probability distribution to predicted class\n",
    "      y_pred = y_pred.argmax(dim=1)\n",
    "\n",
    "      # Compute how many correct predictions we got from batch\n",
    "      for i in range(len(y_pred)):\n",
    "        correct += 1 if y_pred[i] == y[i] else 0\n",
    "\n",
    "    print('\\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format(\n",
    "      loss / len(dataloader.dataset), correct, len(dataloader.dataset),\n",
    "      100. * correct / len(dataloader.dataset)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Putting everything together\n",
    "We now have everything in place to train our model. We define two new hyperparameters to control how many epochs we want to run for and the learning rate. Our loss function will be `nn.CrossEntropyLoss()` and our optimizer will be `optim.SGD(..)`. Note that there are many other (possibly better) options so have a look around!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "\n",
    "N_EPOCHS = 10\n",
    "LEARNING_RATE = 0.1\n",
    "\n",
    "model = Model()\n",
    "\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=LEARNING_RATE)\n",
    "\n",
    "for epoch in range(N_EPOCHS):\n",
    "  train(epoch, model, trainloader, optimizer, loss_fn)\n",
    "  test(model, testloader, loss_fn)\n",
    "\n",
    "# Save model\n",
    "torch.save(model.state_dict(), './model.pt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And that's it! We have now trained a simple model to recognize handwritten digits using PyTorch.\n",
    "Although this is a very simple example, it shows how to go about using PyTorch from start to end and those principles will be applied for more or less any PyTorch project."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extensions\n",
    "\n",
    "Recommended extensions:\n",
    "- Playing around with model and hyperparameters\n",
    "- Create a KMNIST classifier [ https:/rois-codh/kmnist ]\n",
    "- Create a Fashion MNIST classifier [ https:/zalandoresearch/fashion-mnist ]\n",
    "- Look into CNNs for image classification"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_William Profit (williamprofit.com) on behalf of ICDSS (icdss.uk)_"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}