03-modules-and-mlps.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Why Modules\n",
    "\n",
    "A typical training procedure for a neural net:\n",
    "\n",
    "0. Define a dataset ($X$ and $Y$)\n",
    "1. Define the neural network with some learnable weights\n",
    "2. Iterate over the dataset\n",
    "3. Pass inputs to the network (forward pass)\n",
    "4. Compute the loss\n",
    "5. Compute gradients w.r.t. network's weights (backward pass)\n",
    "6. Update weights (e.g., weight = weight - lr * gradient)\n",
    "\n",
    "PyTorch handles 1-6 for you via encapsulation, so you still have the flexibility to change something in between if you want! "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example: MNIST classifier\n",
    "\n",
    "The MNIST dataset is composed of images of digits that must be classified with labels from 0 to 9. The inputs are 28x28 matrices containing the grayscale intensity in each pixel.\n",
    "\n",
    "We will download the MNIST dataset for training a classifier. PyTorch provides a convenient function for that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets\n",
    "import matplotlib.pyplot as plt\n",
    "torch.manual_seed(0);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset\n",
    "It's easy to create your `Dataset`,\n",
    "but PyTorch comes with several built-in datasets for [vision](https://pytorch.org/vision/stable/datasets.html), [audio](https://pytorch.org/audio/stable/datasets.html), and [text](https://pytorch.org/text/stable/datasets.html) modalities.\n",
    "\n",
    "The class `Dataset` gives you information about the number of samples (implement `__len__`) and gives you the sample at a given index (implement `__getitem__`). It's a nice and simple abstraction to work with data. It has the following structure:\n",
    "\n",
    "```python\n",
    "class Dataset(object):\n",
    "    def __getitem__(self, index):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def __len__(self):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def __add__(self, other):\n",
    "        return ConcatDataset([self, other])\n",
    "```\n",
    "\n",
    "For now, let's use MNIST. But feel free to use another `Dataset` as an exercise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download MNIST and store it in \"../data\"\n",
    "# PyTorch.datasets also handles caching for you so you don't have to download the dataset twice\n",
    "train_data = datasets.MNIST('../data', train=True, download=True)\n",
    "test_data = datasets.MNIST('../data', train=False)\n",
    "\n",
    "train_x = train_data.data\n",
    "train_y = train_data.targets\n",
    "test_x = test_data.data\n",
    "test_y = test_data.targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_train_examples = train_x.shape[0]\n",
    "n_test_examples = test_x.shape[0]\n",
    "print('Training instances:', n_train_examples)\n",
    "print('Test instances:', n_test_examples)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the shape of our training data to see how many input features we have:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_x.shape, train_y.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And what the images looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "C = 8\n",
    "fig, axs = plt.subplots(3, C, figsize=(12, 4))\n",
    "for i in range(3):\n",
    "    for j in range(C):\n",
    "        axs[i, j].imshow(train_x[i*C + j], cmap='gray')\n",
    "        axs[i, j].set_axis_off()\n",
    "print(train_y[:24].reshape(3, C))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Formatting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each sample is a 28x28 matrix. But we want to represent them as vectors, since our model (which will be a simple MLP) doesn't take any advantage of the 2D nature of the data.\n",
    "\n",
    "So, we reshape the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_features = 28 * 28\n",
    "train_x_vectors = train_x.view(n_train_examples, num_features)\n",
    "print(train_x_vectors.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we reshape an array (or torch tensor, for that matter), we don't need to specify all dimensions. We can leave one as -1, and it will be automatically determined from the size of the data. This is useful when we don't know a priori the shape of some array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_x_vectors = train_x.view(n_train_examples, -1)\n",
    "test_x_vectors = test_x.view(n_test_examples, -1)\n",
    "\n",
    "print(train_x_vectors.shape, test_x_vectors.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Also, the values are integers in the range $[0, 255]$. It is better to work with float values in a smaller interval, such as $[0, 1]$ or $[-1, 1]$. There are some more elaborate normalization techniques, but for now let's just normalize the data into $[0, 1]$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_x_norm = train_x_vectors / 255.0\n",
    "test_x_norm = test_x_vectors / 255.0\n",
    "print(train_x_norm.max(), train_x_norm.min(), train_x_norm.mean(), train_x_norm.std())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's check all the available labels:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(torch.unique(train_y))\n",
    "num_classes = len(torch.unique(train_y))\n",
    "print('Num classes:', num_classes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modules and MLPs\n",
    "\n",
    "We've seen how the internals of a simple linear classifier work. However, we still had to set a lot of things manually. It's much better to have a higher-level API that encapsulates the classifier.\n",
    "\n",
    "We are going to see that now, with pytorch Module objects. Then, it will allow us to build more complex models, like a multilayer perceptron."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We begin by loading, reshaping and normalizing the data again (so the code looks concise):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.transforms import ToTensor\n",
    "\n",
    "train_dataset = datasets.MNIST('../data', train=True, download=True, transform=ToTensor())\n",
    "test_dataset = datasets.MNIST('../data', train=False, transform=ToTensor())\n",
    "\n",
    "train_x = train_dataset.data\n",
    "train_y = train_dataset.targets\n",
    "test_x = test_dataset.data\n",
    "test_y = test_dataset.targets\n",
    "\n",
    "num_features = 28 * 28\n",
    "num_classes = len(torch.unique(train_y))\n",
    "new_shape = [-1, num_features]\n",
    "train_x_vectors = train_x.reshape(new_shape)\n",
    "test_x_vectors = test_x.reshape(new_shape)\n",
    "\n",
    "# shorten the names\n",
    "train_x = train_x_vectors.float() / 255\n",
    "test_x = test_x_vectors.float() / 255"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Modules\n",
    "\n",
    "PyTorch provides some basic building blocks for neural nets under `.nn` module. Here you can check the complete list of available blocks: https://pytorch.org/docs/stable/nn.html\n",
    "\n",
    "For now, let's recreate a simple linear model using `nn.Linear` (see [doc](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html#torch.nn.Linear))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LinearModel(nn.Module):\n",
    "    def __init__(self, n_features, n_classes):\n",
    "        super().__init__()\n",
    "        self.linear_layer = nn.Linear(n_features, n_classes)\n",
    "        \n",
    "    def forward(self, X):\n",
    "        # This is the same as doing:\n",
    "        # return X @ self.linear_layer.weight.t() + self.linear_layer.bias\n",
    "        # where weight and bias are instances of nn.Parameter\n",
    "        return self.linear_layer(X)\n",
    "\n",
    "linear_model = LinearModel(num_features, num_classes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As before, the model can be called as function in order to produce an output:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = train_x[:2]\n",
    "outputs = linear_model(batch)\n",
    "outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Same as doing the forward method $$w^T x + b$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch @ linear_model.linear_layer.weight.t() + linear_model.linear_layer.bias"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we defined our model, we just have to: \n",
    "- define an iterator\n",
    "- define and compute the loss\n",
    "- compute gradients\n",
    "- define the strategy to update the parameters of our model\n",
    "- glue previous steps to form the training loop!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Batching\n",
    "\n",
    "Batching can be boring to code. PyTorch provides the `DataLoader` class to help us! Dealing with data is one of the most important yet more time consuming tasks. Take a look in the PyTorch `data` submodule to [learn more](https://pytorch.org/docs/stable/data.html).\n",
    "\n",
    "In general, we just have to pass a torch `Dataset` object as input to the dataloader, and then set some hyperparams for the iterator: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "print(type(train_dataset))\n",
    "\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Loss\n",
    "\n",
    "Here is the complete list of available [loss functions](https://pytorch.org/docs/stable/nn.html#loss-functions).\n",
    "If the provided loss functions don't satisfy your constraints, it is easy to define your own loss function: just use torch operations (and be careful with differentiability issues). For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():  # disable gradient-tracking\n",
    "    \n",
    "    dummy_loss = nn.CrossEntropyLoss()\n",
    "    \n",
    "    # try other losses!\n",
    "    # multi-class classification hinge loss (margin-based loss):\n",
    "    # dummy_loss = nn.MultiMarginLoss()  \n",
    "    batch = train_x[:2]\n",
    "    targets = train_y[:2]\n",
    "    predictions = linear_model(batch)\n",
    "    \n",
    "    print(predictions.shape, targets.shape)\n",
    "    print(dummy_loss(predictions, targets))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And writing our own function (from the definition of the Cross Entropy loss):\n",
    "\n",
    "$$\n",
    "CE(p,y) = - \\log\\frac{\\exp(p_y)}{\\sum_c \\exp(p_c)}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dummy_loss(y_pred, y):\n",
    "    one_hot = y.unsqueeze(1) == torch.arange(num_classes).unsqueeze(0)\n",
    "    res = - torch.log(torch.exp(y_pred) / torch.exp(y_pred).sum(-1).unsqueeze(-1))[one_hot]\n",
    "    return res.mean()  # average per sample\n",
    "\n",
    "print(dummy_loss(predictions, targets))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use the CrossEntropy function as our loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_function = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Optimizer\n",
    "\n",
    "The optimizer is the object which handles the update of the model's parameters. In the previous exercise, we were using the famous \"delta\" rule to update our weights:\n",
    "\n",
    "$$\\mathbf{w}_t = \\mathbf{w}_{t-1} - \\alpha \\frac{\\partial L}{\\partial \\mathbf{w}}.$$\n",
    "\n",
    "But there are more ellaborate ways of updating our parameters: \n",
    "\n",
    "<!-- <img src=\"http://cs231n.github.io/assets/nn3/opt2.gif\" width=\"45%\" /> -->\n",
    "\n",
    "<img src=\"http://cs231n.github.io/assets/nn3/opt1.gif\" width=\"45%\" />\n",
    "\n",
    "\n",
    "PyTorch provides an extensive list of optimizers: https://pytorch.org/docs/stable/optim.html. Notice that, as everything else, it should be easy to define your own optimizer procedure. \n",
    "\n",
    "We will use the simple yet powerful SGD optmizer. The optimizer needs to be told which are the parameters to optimize."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters = linear_model.parameters()  # we will optimize all model's parameters!\n",
    "optimizer = torch.optim.SGD(parameters, lr=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Training loop\n",
    "\n",
    "Now we write the main training loop. This is the basic skeleton for training PyTorch models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, dataloader, optimizer, loss_function, num_epochs=1):\n",
    "    # Tell PyTorch that we are in training mode.\n",
    "    # This is useful for mechanisms that work differently during training and test time, like Dropout. \n",
    "    model.train()\n",
    "    \n",
    "    losses = []\n",
    "    for epoch in range(1, num_epochs+1):\n",
    "        print('Starting epoch %d' % epoch)\n",
    "        total_loss = 0\n",
    "        hits = 0\n",
    "\n",
    "        for batch_x, batch_y in dataloader:\n",
    "            # check shapes with:\n",
    "            # import ipdb; ipdb.set_trace()\n",
    "            # batch_x.shape is (batch_size, 28, 28)\n",
    "            # batch_y.shape is (batch_size, )\n",
    "            \n",
    "            # Step 1. Remember that PyTorch accumulates gradients.\n",
    "            # We need to clear them out before each step\n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            # Step 2. Preprocess the data\n",
    "            # (batch_size, 28, 28) -> (batch_size, 784 = 28 * 28)\n",
    "            batch_x = batch_x.reshape(batch_x.shape[0], -1)\n",
    "            batch_x = batch_x.to(torch.float) / 255.0\n",
    "\n",
    "            # Step 3. Run forward pass.\n",
    "            logits = model(batch_x)\n",
    "\n",
    "            # Step 4. Compute loss\n",
    "            loss = loss_function(logits, batch_y)\n",
    "            \n",
    "            # Step 5. Compute gradeints\n",
    "            loss.backward()\n",
    "            \n",
    "            # Step 6. After determining the gradients, take a step toward their (neg-)direction\n",
    "            optimizer.step()\n",
    "            \n",
    "            # Optional. Save statistics of your training\n",
    "            loss_value = loss.item()\n",
    "            total_loss += loss_value\n",
    "            losses.append(loss_value)\n",
    "            y_pred = logits.argmax(dim=1)\n",
    "            hits += torch.sum(y_pred == batch_y).item()\n",
    "        \n",
    "        avg_loss = total_loss / len(train_dataloader.dataset)\n",
    "        print('Epoch loss: %.4f' % avg_loss)\n",
    "        acc = hits / len(train_dataloader.dataset)\n",
    "        print('Epoch accuracy: %.4f' % acc)\n",
    "    \n",
    "    print('Done!')\n",
    "    return losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "linear_losses = train_model(linear_model, train_dataloader, optimizer, loss_function, num_epochs=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Graphics are good to understand the performance of a model. Let's plot the loss curve by training step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(linear_losses, \"-\")\n",
    "ax.set_xlabel('Step')\n",
    "ax.set_ylabel('Loss');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What can you conclude from this?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multilayer Perceptron\n",
    "\n",
    "We can now proceed to a more sofisticated classifier: a multilayer perceptron. Let's build one using the Sequential API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(nn.Module):\n",
    "    def __init__(self, n_features, hidden_size, n_classes):\n",
    "        super().__init__()\n",
    "        linear_layer1 = nn.Linear(n_features, hidden_size)\n",
    "        linear_layer2 = nn.Linear(hidden_size, hidden_size)\n",
    "        linear_layer3 = nn.Linear(hidden_size, n_classes)\n",
    "        self.feedforward = nn.Sequential(\n",
    "            linear_layer1, \n",
    "            nn.Tanh(), \n",
    "            linear_layer2, \n",
    "            nn.Tanh(),\n",
    "            linear_layer3\n",
    "        )\n",
    "\n",
    "    def forward(self, X):\n",
    "        return self.feedforward(X)\n",
    "\n",
    "hidden_size = 200\n",
    "mlp = MLP(num_features, hidden_size, num_classes)\n",
    "loss_function = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(mlp.parameters(), lr=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's train the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlp_losses = train_model(mlp, train_dataloader, optimizer, loss_function, num_epochs=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How do the loss and accuracy compare with the linear model?\n",
    "\n",
    "You probably also noticed a difference in running time!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(linear_losses, \".\", label=\"linear\")\n",
    "ax.plot(mlp_losses, \".\", label=\"mlp\")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the different concentration of dots in the MLP and Linear graphics!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Validation data\n",
    "\n",
    "Evaluating the performance on training data is important to understand if the model is actually learning, but if we want to know if our model has any usefulness, we should evaluate its performance on validation or test data.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_model(model, test_x, test_y):\n",
    "    # Tell PyTorch that we are in evaluation mode.\n",
    "    model.eval()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        loss_function = torch.nn.CrossEntropyLoss()\n",
    "        logits = model(test_x)\n",
    "        loss = loss_function(logits, test_y)\n",
    "\n",
    "        y_pred = logits.argmax(dim=1)\n",
    "        hits = torch.sum(y_pred == test_y).item()\n",
    "    \n",
    "    return loss.item() / len(test_x), hits / len(test_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate_model(mlp, train_x, train_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate_model(mlp, test_x, test_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate_model(linear_model, train_x, train_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate_model(linear_model, test_x, test_y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How can we make our model better? There are two things to be done:\n",
    "\n",
    "1. **Hyperparameter search**. Do a grid search or random search on the hyperparameters (hidden size, learning rate, batch size, activation function, type of optimizer, ...)\n",
    "2. **Generalize better**. This include either finding some better feature representation or regularizing, i.e., add some kind of penalty to the model weights that encourages it to find a more general solution. Examples: L2-norm weight regularization, dropout.\n",
    "3. **Early stop**. Evaluate the model on validation data after each epoch or some number of batches; only save it when validation performance increases. This means detecting when the model achieved its performance peak."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Dropout\n",
    "\n",
    "We could try dropout. It effectivelly deactivates some neural connections at random, forcing the network to avoid depending on specific inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLPDropout(nn.Module):\n",
    "    def __init__(self, n_features, hidden_size, n_classes, p_dropout):\n",
    "        super().__init__()\n",
    "        linear_layer1 = nn.Linear(n_features, hidden_size)\n",
    "        linear_layer2 = nn.Linear(hidden_size, n_classes)\n",
    "        self.feedforward = nn.Sequential(\n",
    "            linear_layer1,\n",
    "            nn.Tanh(),\n",
    "            nn.Dropout(p_dropout),\n",
    "            linear_layer2\n",
    "        )\n",
    "\n",
    "    def forward(self, X):\n",
    "        return self.feedforward(X)\n",
    "\n",
    "hidden_size = 200\n",
    "p_dropout = 0.5\n",
    "mlp_dropout = MLPDropout(num_features, hidden_size, num_classes, p_dropout)\n",
    "loss_function = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(mlp_dropout.parameters(), lr=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "losses = train_model(mlp_dropout, train_dataloader, optimizer, loss_function, num_epochs=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training loss is a bit worse, as expected. After all, we are obstructing some connections.\n",
    "\n",
    "Now let's check validation performance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate_model(mlp, test_x, test_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate_model(mlp_dropout, test_x, test_y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "No improvement. Ideally, we should retrain our model with different hyperparamters (learning rates, layer sizes, number of layers, dropout rate) as well as some changes in the structure (different optimizers, activation functions, losses). However, data representation plays a key role. \n",
    "\n",
    "<br>\n",
    "<center>\n",
    "<i>Do you think representing the input as independent pixels is a good idea for recognizing digits?</i>\n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Saving\n",
    "\n",
    "Persisting the model after training is obviously important to reuse it later. In Pytorch, we can save the model calling `save()` and passing  the model's `state_dict` (a Python dict that maps all parameters name to their actual tensors)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(mlp.state_dict(), 'mlp.model')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Later, recreate the model and load the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlp2 = MLP(num_features, hidden_size, num_classes)\n",
    "mlp2.load_state_dict(torch.load('mlp.model'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check the performance to see if it's the same!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate_model(mlp, test_x, test_y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The End\n",
    "\n",
    "![https://twitter.com/karpathy/status/1013244313327681536](img/common_mistakes.png)\n",
    "https://twitter.com/karpathy/status/1013244313327681536"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercises\n",
    "\n",
    "- Try running the MLP example for more epochs\n",
    "- Try using CNNs: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}