introduction_pytorch_poutyne.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to PyTorch and Poutyne\n",
    "\n",
    "In this notebook, we train a simple fully-connected network and a simple convolutional network on MNIST. First, we train it by coding our own training loop as the PyTorch library expects of us to. Then, we use Poutyne to simplify our code.\n",
    "\n",
    "Let's install the latest version of Poutyne (if it's not already) and import all the needed packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --upgrade poutyne\n",
    "%pip install --upgrade colorama # Uncomment this line to enable colors.\n",
    "%matplotlib inline\n",
    "import os\n",
    "import math\n",
    "import shutil\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import random_split, DataLoader\n",
    "from torchvision import transforms, utils\n",
    "from torchvision.datasets.mnist import MNIST\n",
    "\n",
    "from poutyne import set_seeds, Model, ModelCheckpoint, CSVLogger, ModelBundle"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Also, we need to set Python's, NumPy's and PyTorch's seeds by using Poutyne function so that our training is (almost) reproducible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_seeds(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basis of Training a Neural Network\n",
    "\n",
    "In **stochastic gradient descent**, a **batch** of `m` examples are drawn from the train dataset. In the so-called forward pass, these examples are passed through the neural network and an average of their loss values is done. In the backward pass, the average loss is backpropagated through the network to compute the gradient of each parameter. In practice, the `m` examples of a batch are drawn without replacement. Thus, we define one **epoch** of training being the number of batches needed to loop through the entire training dataset.\n",
    "\n",
    "In addition to the training dataset, a **validation dataset** is used to evaluate the neural network at the end of each epoch. This validation dataset can be used to select the best model during training and thus avoiding overfitting the training set. It also can have other uses such as selecting hyperparameters\n",
    "\n",
    "Finally, a **test dataset** is used at the end to evaluate the final model.\n",
    "\n",
    "### Training constants"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now, let's set our training constants. We first have the CUDA device used for training if one is present. Second, we set the train_split to 0.8 (80%) to use 80% of the dataset for training and 20% for testing the trained model. Third, we set the number of classes (i.e. one for each number). Finally, we set the batch size (i.e. the number of elements to see before updating the model), the learning rate for the optimizer, and the number of epochs (i.e. the number of times we see the full dataset).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cuda_device = 0\n",
    "device = torch.device(\"cuda:%d\" % cuda_device if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "train_split_percent = 0.8\n",
    "\n",
    "num_classes = 10\n",
    "\n",
    "batch_size = 32\n",
    "learning_rate = 0.1\n",
    "num_epochs = 5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading the MNIST dataset\n",
    "\n",
    "The following code helps load the MNIST dataset and creates the PyTorch DataLoaders that split our datasets into batches. The train DataLoader shuffles the examples of the train dataset to draw the examples without replacement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "full_train_dataset = MNIST(\n",
    "    \"./datasets/\", train=True, download=True, transform=transforms.ToTensor()\n",
    ")\n",
    "test_dataset = MNIST(\n",
    "    \"./datasets/\", train=False, download=True, transform=transforms.ToTensor()\n",
    ")\n",
    "\n",
    "num_data = len(full_train_dataset)\n",
    "train_length = int(math.floor(train_split_percent * num_data))\n",
    "valid_length = num_data - train_length\n",
    "\n",
    "train_dataset, valid_dataset = random_split(\n",
    "    full_train_dataset,\n",
    "    [train_length, valid_length],\n",
    "    generator=torch.Generator().manual_seed(42),\n",
    ")\n",
    "\n",
    "train_loader = DataLoader(\n",
    "    train_dataset, batch_size=batch_size, num_workers=2, shuffle=True\n",
    ")\n",
    "valid_loader = DataLoader(valid_dataset, batch_size=batch_size, num_workers=2)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size, num_workers=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at some examples of the dataset by looking at the first batch in our train DataLoader and formatting it into a grid and plotting it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "inputs = next(iter(train_loader))[0]\n",
    "input_grid = utils.make_grid(inputs)\n",
    "\n",
    "fig = plt.figure(figsize=(10, 10))\n",
    "inp = input_grid.numpy().transpose((1, 2, 0))\n",
    "plt.imshow(inp)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Neural Network Architectures\n",
    "\n",
    "We train a fully-connected neural network and a convolutional neural network with approximately the same number of parameters.\n",
    "\n",
    "#### Fully-connected Network\n",
    "In short, the fully-connected network follows this architecture: ``Input -> [Linear -> ReLU]*3 -> Linear``. The following table shows it in details:\n",
    "\n",
    "| Layer Type                  | Output Size |       # of Parameters       |\n",
    "|-----------------------------|:-----------:|:---------------------------:|\n",
    "| Input                       |   1x28x28   |              0              |\n",
    "| Flatten                     |  1\\*28\\*28  |              0              |\n",
    "| **Linear with 256 neurons** |     256     | 28\\*28\\*256 + 256 = 200,960 |\n",
    "| ReLU                        |      *      |              0              |\n",
    "| **Linear with 128 neurons** |     128     |   256*128 + 128 = 32,896    |\n",
    "| ReLU                        |      *      |              0              |\n",
    "| **Linear with 64 neurons**  |     64      |     128*64 + 64 = 8,256     |\n",
    "| ReLU                        |      *      |              0              |\n",
    "| **Linear with 10 neurons**  |      10     |      64*10 + 10 = 650       |\n",
    "\n",
    "Total # of parameters of the fully-connected network: 242,762\n",
    "\n",
    "#### Convolutional Network\n",
    "\n",
    "The convolutional neural network architecture starts with some convolution and max-pooling layers. These are then followed by fully-connected layers. We calculate the total number of parameters that the network needs. In short, the convolutional network follows this architecture: ``Input -> [Conv -> ReLU -> MaxPool]*2 -> Dropout -> Linear -> ReLU -> Dropout -> Linear``. The following table shows it in details:\n",
    "\n",
    "| Layer Type                                     | Output Size |        # of Parameters        |\n",
    "|------------------------------------------------|:-----------:|:-----------------------------:|\n",
    "| Input                                          |   1x28x28   |               0               |\n",
    "| **Conv with 16 3x3 filters with padding of 1** |   16x28x28  |    16\\*1\\*3\\*3 + 16 = 160     |\n",
    "| ReLU                                           |   16x28x28  |               0               |\n",
    "| MaxPool 2x2                                    |   16x14x14  |               0               |\n",
    "| **Conv with 32 3x3 filters with padding of 1** |   32x14x14  |   32\\*16\\*3\\*3 + 32 = 4,640   |\n",
    "| ReLU                                           |   32x14x14  |               0               |\n",
    "| MaxPool 2x2                                    |    32x7x7   |               0               |\n",
    "| Dropout of 0.25                                |    32x7x7   |               0               |\n",
    "| Flatten                                        |   32\\*7\\*7  |               0               |\n",
    "| **Linear with 128 neurons**                    |     128     | 32\\*7\\*7\\*128 + 128 = 200,832 |\n",
    "| ReLU                                           |     128     |               0               |\n",
    "| Dropout of 0.5                                 |     128     |               0               |\n",
    "| **Linear with 10 neurons**                     |      10     |        128\\*10 + 10 = 1290    |\n",
    "\n",
    "Total # of parameters of the convolutional network: 206,922"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def create_fully_connected_network():\n",
    "    \"\"\"\n",
    "    This function returns the fully-connected network layed out above.\n",
    "    \"\"\"\n",
    "    return nn.Sequential(\n",
    "        nn.Flatten(),\n",
    "        nn.Linear(28 * 28, 256),\n",
    "        nn.ReLU(),\n",
    "        nn.Linear(256, 128),\n",
    "        nn.ReLU(),\n",
    "        nn.Linear(128, 64),\n",
    "        nn.ReLU(),\n",
    "        nn.Linear(64, num_classes),\n",
    "    )\n",
    "\n",
    "\n",
    "def create_convolutional_network():\n",
    "    \"\"\"\n",
    "    This function returns the convolutional network layed out above.\n",
    "    \"\"\"\n",
    "    return nn.Sequential(\n",
    "        nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, padding=1),\n",
    "        nn.ReLU(),\n",
    "        nn.MaxPool2d(2),\n",
    "        nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1),\n",
    "        nn.ReLU(),\n",
    "        nn.MaxPool2d(2),\n",
    "        nn.Dropout(0.25),\n",
    "        nn.Flatten(),\n",
    "        nn.Linear(32 * 7 * 7, 128),\n",
    "        nn.ReLU(),\n",
    "        nn.Dropout(0.5),\n",
    "        nn.Linear(128, num_classes),\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training the PyTorch way \n",
    "\n",
    "That is, writing your own training loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pytorch_accuracy(y_pred, y_true):\n",
    "    \"\"\"\n",
    "    Computes the accuracy for a batch of predictions\n",
    "\n",
    "    Args:\n",
    "        y_pred (torch.Tensor): the logit predictions of the neural network.\n",
    "        y_true (torch.Tensor): the ground truths.\n",
    "\n",
    "    Returns:\n",
    "        The average accuracy of the batch.\n",
    "    \"\"\"\n",
    "    y_pred = y_pred.argmax(1)\n",
    "    return (y_pred == y_true).float().mean() * 100\n",
    "\n",
    "\n",
    "def pytorch_train_one_epoch(network, optimizer, loss_function):\n",
    "    \"\"\"\n",
    "    Trains the neural network for one epoch on the train DataLoader.\n",
    "\n",
    "    Args:\n",
    "        network (torch.nn.Module): The neural network to train.\n",
    "        optimizer (torch.optim.Optimizer): The optimizer of the neural network\n",
    "        loss_function: The loss function.\n",
    "\n",
    "    Returns:\n",
    "        A tuple (loss, accuracy) corresponding to an average of the losses and\n",
    "        an average of the accuracy, respectively, on the train DataLoader.\n",
    "    \"\"\"\n",
    "    network.train(True)\n",
    "    with torch.enable_grad():\n",
    "        loss_sum = 0.0\n",
    "        acc_sum = 0.0\n",
    "        example_count = 0\n",
    "        for (x, y) in train_loader:\n",
    "            # Transfer batch on GPU if needed.\n",
    "            x = x.to(device)\n",
    "            y = y.to(device)\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            y_pred = network(x)\n",
    "\n",
    "            loss = loss_function(y_pred, y)\n",
    "\n",
    "            loss.backward()\n",
    "\n",
    "            optimizer.step()\n",
    "\n",
    "            # Since the loss and accuracy are averages for the batch, we multiply\n",
    "            # them by the the number of examples so that we can do the right\n",
    "            # averages at the end of the epoch.\n",
    "            loss_sum += float(loss) * len(x)\n",
    "            acc_sum += float(pytorch_accuracy(y_pred, y)) * len(x)\n",
    "            example_count += len(x)\n",
    "\n",
    "    avg_loss = loss_sum / example_count\n",
    "    avg_acc = acc_sum / example_count\n",
    "    return avg_loss, avg_acc\n",
    "\n",
    "\n",
    "def pytorch_test(network, loader, loss_function):\n",
    "    \"\"\"\n",
    "    Tests the neural network on a DataLoader.\n",
    "\n",
    "    Args:\n",
    "        network (torch.nn.Module): The neural network to test.\n",
    "        loader (torch.utils.data.DataLoader): The DataLoader to test on.\n",
    "        loss_function: The loss function.\n",
    "\n",
    "    Returns:\n",
    "        A tuple (loss, accuracy) corresponding to an average of the losses and\n",
    "        an average of the accuracy, respectively, on the DataLoader.\n",
    "    \"\"\"\n",
    "    network.eval()\n",
    "    with torch.no_grad():\n",
    "        loss_sum = 0.0\n",
    "        acc_sum = 0.0\n",
    "        example_count = 0\n",
    "        for (x, y) in loader:\n",
    "            # Transfer batch on GPU if needed.\n",
    "            x = x.to(device)\n",
    "            y = y.to(device)\n",
    "\n",
    "            y_pred = network(x)\n",
    "            loss = loss_function(y_pred, y)\n",
    "\n",
    "            # Since the loss and accuracy are averages for the batch, we multiply\n",
    "            # them by the the number of examples so that we can do the right\n",
    "            # averages at the end of the test.\n",
    "            loss_sum += float(loss) * len(x)\n",
    "            acc_sum += float(pytorch_accuracy(y_pred, y)) * len(x)\n",
    "            example_count += len(x)\n",
    "\n",
    "    avg_loss = loss_sum / example_count\n",
    "    avg_acc = acc_sum / example_count\n",
    "    return avg_loss, avg_acc\n",
    "\n",
    "\n",
    "def pytorch_train(network):\n",
    "    \"\"\"\n",
    "    This function transfers the neural network to the right device,\n",
    "    trains it for a certain number of epochs, tests at each epoch on\n",
    "    the validation set and outputs the results on the test set at the\n",
    "    end of training.\n",
    "\n",
    "    Args:\n",
    "        network (torch.nn.Module): The neural network to train.\n",
    "\n",
    "    Example:\n",
    "        This function displays something like this:\n",
    "\n",
    "        .. code-block:: python\n",
    "\n",
    "            Epoch 1/5: loss: 0.5026924496193726, acc: 84.26666259765625, val_loss: 0.17258917854229608, val_acc: 94.75\n",
    "            Epoch 2/5: loss: 0.13690324830015502, acc: 95.73332977294922, val_loss: 0.14024296019474666, val_acc: 95.68333435058594\n",
    "            Epoch 3/5: loss: 0.08836929737279813, acc: 97.29582977294922, val_loss: 0.10380942322810491, val_acc: 96.66666412353516\n",
    "            Epoch 4/5: loss: 0.06714504160980383, acc: 97.91874694824219, val_loss: 0.09626663728555043, val_acc: 97.18333435058594\n",
    "            Epoch 5/5: loss: 0.05063822727650404, acc: 98.42708587646484, val_loss: 0.10017542181412378, val_acc: 96.95833587646484\n",
    "            Test:\n",
    "                Loss: 0.09501855444908142\n",
    "                Accuracy: 97.12999725341797\n",
    "    \"\"\"\n",
    "    print(network)\n",
    "\n",
    "    # Transfer weights on GPU if needed.\n",
    "    network.to(device)\n",
    "\n",
    "    optimizer = optim.SGD(network.parameters(), lr=learning_rate)\n",
    "    loss_function = nn.CrossEntropyLoss()\n",
    "\n",
    "    for epoch in range(1, num_epochs + 1):\n",
    "        # Training the neural network via backpropagation\n",
    "        train_loss, train_acc = pytorch_train_one_epoch(\n",
    "            network, optimizer, loss_function\n",
    "        )\n",
    "\n",
    "        # Validation at the end of the epoch\n",
    "        valid_loss, valid_acc = pytorch_test(network, valid_loader, loss_function)\n",
    "\n",
    "        print(\n",
    "            \"Epoch {}/{}: loss: {}, acc: {}, val_loss: {}, val_acc: {}\".format(\n",
    "                epoch, num_epochs, train_loss, train_acc, valid_loss, valid_acc\n",
    "            )\n",
    "        )\n",
    "\n",
    "    # Test at the end of the training\n",
    "    test_loss, test_acc = pytorch_test(network, test_loader, loss_function)\n",
    "    print(\"Test:\\n\\tLoss: {}\\n\\tAccuracy: {}\".format(test_loss, test_acc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fc_net = create_fully_connected_network()\n",
    "pytorch_train(fc_net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conv_net = create_convolutional_network()\n",
    "pytorch_train(conv_net)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training the Poutyne way\n",
    "\n",
    "That is, only few lines of code with a better output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def poutyne_train(network):\n",
    "    \"\"\"\n",
    "    This function creates a Poutyne Model (see https://poutyne.org/model.html), sends the Model\n",
    "    on the specified device, and uses the `fit_generator` method to train the neural network.\n",
    "    At the end, the `evaluate_generator` is used on  the test set.\n",
    "\n",
    "    Args:\n",
    "        network (torch.nn.Module): The neural network to train.\n",
    "    \"\"\"\n",
    "    print(network)\n",
    "\n",
    "    optimizer = optim.SGD(network.parameters(), lr=learning_rate)\n",
    "    loss_function = nn.CrossEntropyLoss()\n",
    "\n",
    "    # Poutyne Model on GPU\n",
    "    model = Model(\n",
    "        network,\n",
    "        optimizer,\n",
    "        loss_function,\n",
    "        batch_metrics=[\"accuracy\"],\n",
    "        device=device,\n",
    "    )\n",
    "\n",
    "    # Train\n",
    "    model.fit_generator(train_loader, valid_loader, epochs=num_epochs)\n",
    "\n",
    "    # Test\n",
    "    test_loss, test_acc = model.evaluate_generator(test_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fc_net = create_fully_connected_network()\n",
    "poutyne_train(fc_net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "conv_net = create_convolutional_network()\n",
    "poutyne_train(conv_net)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Poutyne Callbacks\n",
    "\n",
    "One nice feature of Poutyne is [callbacks](https://poutyne.org/callbacks.html). Callbacks allow to do actions during training of the neural network. In the following example, we use 3 callbacks. One that saves the latest weights in a file to be able to continue the optimization at the end of training if more epochs are needed. Another one that saves the best weights according to the performance on the validation dataset. Finally, another one that saves the displayed logs into a TSV file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_with_callbacks(name, network):\n",
    "    \"\"\"\n",
    "    In addition to the the `poutyne_train`, this function saves checkpoints and logs as described above.\n",
    "\n",
    "    Args:\n",
    "        name (str): a name used to save logs and checkpoints.\n",
    "        network (torch.nn.Module): The neural network to train.\n",
    "    \"\"\"\n",
    "    print(network)\n",
    "\n",
    "    # We are saving everything into ./saves/{name}.\n",
    "    save_path = os.path.join(\"saves\", name)\n",
    "\n",
    "    # Creating saving directory if necessary.\n",
    "    os.makedirs(save_path, exist_ok=True)\n",
    "\n",
    "    callbacks = [\n",
    "        # Save the latest weights to be able to continue the optimization at the end for more epochs.\n",
    "        ModelCheckpoint(os.path.join(save_path, \"last_epoch.ckpt\")),\n",
    "        # Save the weights in a new file when the current model is better than all previous models.\n",
    "        ModelCheckpoint(\n",
    "            os.path.join(save_path, \"best_epoch_{epoch}.ckpt\"),\n",
    "            monitor=\"val_acc\",\n",
    "            mode=\"max\",\n",
    "            save_best_only=True,\n",
    "            restore_best=True,\n",
    "            verbose=True,\n",
    "        ),\n",
    "        # Save the losses and accuracies for each epoch in a TSV.\n",
    "        CSVLogger(os.path.join(save_path, \"log.tsv\"), separator=\"\\t\"),\n",
    "    ]\n",
    "\n",
    "    optimizer = optim.SGD(network.parameters(), lr=learning_rate)\n",
    "    loss_function = nn.CrossEntropyLoss()\n",
    "\n",
    "    model = Model(\n",
    "        network,\n",
    "        optimizer,\n",
    "        loss_function,\n",
    "        batch_metrics=[\"accuracy\"],\n",
    "        device=device,\n",
    "    )\n",
    "    model.fit_generator(\n",
    "        train_loader, valid_loader, epochs=num_epochs, callbacks=callbacks\n",
    "    )\n",
    "\n",
    "    test_loss, test_acc = model.evaluate_generator(test_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fc_net = create_fully_connected_network()\n",
    "train_with_callbacks(\"fc\", fc_net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conv_net = create_convolutional_network()\n",
    "train_with_callbacks(\"conv\", conv_net)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Making Your Own Callback\n",
    "\n",
    "While Poutyne provides a great number of [predefined callbacks](https://poutyne.org/callbacks.html), it is sometimes useful to make your own callback. In addition to the [documentation](https://poutyne.org/callbacks.html#poutyne.framework.callbacks.Callback) of the Callback class, see the [Tips and Tricks notebook](https://github.com/GRAAL-Research/poutyne/blob/master/examples/tips_and_tricks.ipynb) for an example."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Poutyne ModelBundle\n",
    "\n",
    "Most of the time when using Poutyne (or even Pytorch in general), we will find ourselves in an iterative model hyperparameters finetuning loop. For efficient model search, we will usually wish to save our best performing models, their training and testing statistics and even sometimes wish to retrain an already trained model for further tuning. All of the above can be easily implemented with the flexibility of Poutyne Callbacks, but having to define and initialize each and every Callback object we wish for our model quickly feels cumbersome.\n",
    "\n",
    "This is why Poutyne provides a [ModelBundle class](https://poutyne.org/experiment.html#poutyne.ModelBundle), which aims specifically at enabling quick model iteration search, while not sacrifying on the quality of a single experiment - statistics logging, best models saving, etc. ModelBundle is actually a simple wrapper between a PyTorch network and Poutyne's core Callback objects for logging and saving. Given a working directory where to output the various logging files and a PyTorch network, the ModelBundle class reduces the whole training loop to a single line.\n",
    "\n",
    "The following code uses [Poutyne's ModelBundle class](https://poutyne.org/experiment.html#poutyne.ModelBundle) to train a network for 5 epochs. The code is quite simpler than the code in the Poutyne Callbacks section while doing more (only 3 lines). Once trained for 5 epochs, it is then possible to resume the optimization at the 5th epoch for 5 more epochs until the 10th epoch using the same function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model_bundle(network, name, epochs=5):\n",
    "    \"\"\"\n",
    "    This function creates a Poutyne ModelBundle, trains the input module\n",
    "    on the train loader and then tests its performance on the test loader.\n",
    "    All training and testing statistics are saved, as well as best model\n",
    "    checkpoints.\n",
    "\n",
    "    Args:\n",
    "        network (torch.nn.Module): The neural network to train.\n",
    "        working_directory (str): The directory where to output files to save.\n",
    "        epochs (int): The number of epochs. (Default: 5)\n",
    "    \"\"\"\n",
    "    print(network)\n",
    "\n",
    "    optimizer = optim.SGD(network.parameters(), lr=learning_rate)\n",
    "\n",
    "    # Everything is going to be saved in ./saves/{name}.\n",
    "    save_path = os.path.join(\"saves\", name)\n",
    "\n",
    "    # Poutyne ModelBundle\n",
    "    model_bundle = ModelBundle.from_network(\n",
    "        save_path,\n",
    "        network,\n",
    "        device=device,\n",
    "        optimizer=optimizer,\n",
    "        task=\"classif\",\n",
    "    )\n",
    "\n",
    "    # Train\n",
    "    model_bundle.train(train_loader, valid_loader, epochs=epochs)\n",
    "\n",
    "    # Test\n",
    "    model_bundle.test(test_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first remove the output directory if it exists so that we start training from scratch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shutil.rmtree(\"saves/conv_net_model_bundle\", ignore_errors=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's train for 5 epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conv_net = create_convolutional_network()\n",
    "train_model_bundle(conv_net, \"conv_net_model_bundle\", epochs=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "conv_net = create_convolutional_network()\n",
    "train_model_bundle(conv_net, \"conv_net_model_bundle\", epochs=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next Step\n",
    "We have plenty more examples for you, see [this](https://github.com/GRAAL-Research/poutyne/blob/master/examples/tips_and_tricks.ipynb) notebook for tips and tricks using Poutyne, this [one](https://github.com/GRAAL-Research/poutyne/blob/master/examples/policy_cifar_example.ipynb) and this [one](https://github.com/GRAAL-Research/poutyne/blob/master/examples/policy_interface.ipynb) for policy examples or this [one](https://github.com/GRAAL-Research/poutyne/blob/master/examples/transfer_learning.ipynb) for a transfer learning example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}