05_pytorch_going_modular_script_mode.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Lab-of-Infinity/PyTorch-for-Deep-Learning-Zero-to-Mastery/blob/main/05_pytorch_going_modular_script_mode.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1M5tsObEa17-",
        "tags": []
      },
      "source": [
        "# 05. Going Modular: script mode\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "tags": [],
        "id": "W9w6CQZ5V5-7"
      },
      "source": [
        "## What is script mode?\n",
        "\n",
        "**Script mode** uses [Jupyter Notebook cell magic](https://ipython.readthedocs.io/en/stable/interactive/magics.html) (special commands) to turn specific cells into Python scripts.\n",
        "\n",
        "For example if you run the following code in a cell, you'll create a Python file called `hello_world.py`:\n",
        "\n",
        "```\n",
        "%%writefile hello_world.py\n",
        "print(\"hello world, machine learning is fun!\")\n",
        "```\n",
        "\n",
        "You could then run this Python file on the command line with:\n",
        "\n",
        "```\n",
        "python hello_world.py\n",
        "\n",
        ">>> hello world, machine learning is fun!\n",
        "```\n",
        "\n",
        "The main cell magic we're interested in using is `%%writefile`.\n",
        "\n",
        "Putting `%%writefile filename` at the top of a cell in Jupyter or Google Colab will write the contents of that cell to a specified `filename`.\n",
        "\n",
        "> **Question:** Do I have to create Python files like this? Can't I just start directly with a Python file and skip using a Google Colab notebook?\n",
        ">\n",
        "> **Answer:** Yes. This is only *one* way of creating Python scripts. If you know the kind of script you'd like to write, you could start writing it straight away. But since using Jupyter/Google Colab notebooks is a popular way of starting off data science and machine learning projects, knowing about the `%%writefile` magic command is a handy tip. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "tags": [],
        "id": "TkLUvENvV5_A"
      },
      "source": [
        "## What we're going to cover\n",
        "\n",
        "By the end of this notebook you should finish with a directory structure of: \n",
        "\n",
        "```\n",
        "going_modular/\n",
        "├── going_modular/\n",
        "│   ├── data_setup.py\n",
        "│   ├── engine.py\n",
        "│   ├── model_builder.py\n",
        "│   ├── train.py\n",
        "│   └── utils.py\n",
        "├── models/\n",
        "│   ├── 05_going_modular_cell_mode_tinyvgg_model.pth\n",
        "│   └── 05_going_modular_script_mode_tinyvgg_model.pth\n",
        "└── data/\n",
        "    └── pizza_steak_sushi/\n",
        "        ├── train/\n",
        "        │   ├── pizza/\n",
        "        │   │   ├── image01.jpeg\n",
        "        │   │   └── ...\n",
        "        │   ├── steak/\n",
        "        │   └── sushi/\n",
        "        └── test/\n",
        "            ├── pizza/\n",
        "            ├── steak/\n",
        "            └── sushi/\n",
        "```\n",
        "\n",
        "Using this directory structure, you should be able to train a model from within a notebook with the command:\n",
        "\n",
        "```\n",
        "!python going_modular/train.py\n",
        "```\n",
        "\n",
        "Or from the command line with:\n",
        "\n",
        "```\n",
        "python going_modular/train.py\n",
        "```\n",
        "\n",
        "In essence, we will have turned our helpful notebook code into **reusable modular code**."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OW_3D_fOap2I"
      },
      "source": [
        "## 0. Creating a folder for storing Python scripts\n",
        "\n",
        "Since we're going to be creating Python scripts out of our most useful code cells, let's create a folder for storing those scripts.\n",
        "\n",
        "We'll call the folder `going_modular` and create it using Python's [`os.makedirs()`](https://docs.python.org/3/library/os.html) method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K29qwItveNm9"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "os.makedirs(\"going_modular\", exist_ok=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1XgXFr97bCsI"
      },
      "source": [
        "## 1. Get data\n",
        "\n",
        "We're going to start by downloading the same data we used in `notebook 04`, the `pizza_steak_sushi` dataset with images of pizza, steak and sushi."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wl_jfPzVbGDS",
        "outputId": "51a3e5e8-991c-4309-b1ba-b24d32bbdb61"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "data/pizza_steak_sushi directory exists.\n",
            "Downloading pizza, steak, sushi data...\n",
            "Unzipping pizza, steak, sushi data...\n"
          ]
        }
      ],
      "source": [
        "import os\n",
        "import zipfile\n",
        "\n",
        "from pathlib import Path\n",
        "\n",
        "import requests\n",
        "\n",
        "# Setup path to data folder\n",
        "data_path = Path(\"data/\")\n",
        "image_path = data_path / \"pizza_steak_sushi\"\n",
        "\n",
        "# If the image folder doesn't exist, download it and prepare it... \n",
        "if image_path.is_dir():\n",
        "    print(f\"{image_path} directory exists.\")\n",
        "else:\n",
        "    print(f\"Did not find {image_path} directory, creating one...\")\n",
        "    image_path.mkdir(parents=True, exist_ok=True)\n",
        "    \n",
        "# Download pizza, steak, sushi data\n",
        "with open(data_path / \"pizza_steak_sushi.zip\", \"wb\") as f:\n",
        "    request = requests.get(\"https://github.com/mrdbourke/pytorch-deep-learning/raw/main/data/pizza_steak_sushi.zip\")\n",
        "    print(\"Downloading pizza, steak, sushi data...\")\n",
        "    f.write(request.content)\n",
        "\n",
        "# Unzip pizza, steak, sushi data\n",
        "with zipfile.ZipFile(data_path / \"pizza_steak_sushi.zip\", \"r\") as zip_ref:\n",
        "    print(\"Unzipping pizza, steak, sushi data...\") \n",
        "    zip_ref.extractall(image_path)\n",
        "\n",
        "# Remove zip file\n",
        "os.remove(data_path / \"pizza_steak_sushi.zip\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8qYMdiBMbde0",
        "outputId": "a5093c72-455f-471c-9e95-5b391eea9ebf"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(PosixPath('data/pizza_steak_sushi/train'),\n",
              " PosixPath('data/pizza_steak_sushi/test'))"
            ]
          },
          "execution_count": 3,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Setup train and testing paths\n",
        "train_dir = image_path / \"train\"\n",
        "test_dir = image_path / \"test\"\n",
        "\n",
        "train_dir, test_dir"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eYL6NynybU9Z"
      },
      "source": [
        "## 2. Create Datasets and DataLoaders\n",
        "\n",
        "Let's turn our data into PyTorch `Dataset`'s and `DataLoader`'s and find out a few useful attributes from them such as `classes` and their lengths. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DaOJOd-QbOQk",
        "outputId": "dc7168b3-8bee-40c4-867e-58aa058175b5"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Train data:\n",
            "Dataset ImageFolder\n",
            "    Number of datapoints: 225\n",
            "    Root location: data/pizza_steak_sushi/train\n",
            "    StandardTransform\n",
            "Transform: Compose(\n",
            "               Resize(size=(64, 64), interpolation=bilinear, max_size=None, antialias=None)\n",
            "               ToTensor()\n",
            "           )\n",
            "Test data:\n",
            "Dataset ImageFolder\n",
            "    Number of datapoints: 75\n",
            "    Root location: data/pizza_steak_sushi/test\n",
            "    StandardTransform\n",
            "Transform: Compose(\n",
            "               Resize(size=(64, 64), interpolation=bilinear, max_size=None, antialias=None)\n",
            "               ToTensor()\n",
            "           )\n"
          ]
        }
      ],
      "source": [
        "from torchvision import datasets, transforms\n",
        "\n",
        "# Create simple transform\n",
        "data_transform = transforms.Compose([ \n",
        "    transforms.Resize((64, 64)),\n",
        "    transforms.ToTensor(),\n",
        "])\n",
        "\n",
        "# Use ImageFolder to create dataset(s)\n",
        "train_data = datasets.ImageFolder(root=train_dir, # target folder of images\n",
        "                                  transform=data_transform, # transforms to perform on data (images)\n",
        "                                  target_transform=None) # transforms to perform on labels (if necessary)\n",
        "\n",
        "test_data = datasets.ImageFolder(root=test_dir, \n",
        "                                 transform=data_transform)\n",
        "\n",
        "print(f\"Train data:\\n{train_data}\\nTest data:\\n{test_data}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eEln0VFmbfyY",
        "outputId": "c514574e-e3ab-4689-a5ee-6a2aef6a11d2"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['pizza', 'steak', 'sushi']"
            ]
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Get class names as a list\n",
        "class_names = train_data.classes\n",
        "class_names"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IMGTFVr8bgwE",
        "outputId": "4268d11f-8899-447a-baba-7e91ade01f13"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'pizza': 0, 'steak': 1, 'sushi': 2}"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Can also get class names as a dict\n",
        "class_dict = train_data.class_to_idx\n",
        "class_dict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6a8IW-IYbhGX",
        "outputId": "87bc60ac-3be3-499f-e0a8-aed5964a3075"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(225, 75)"
            ]
          },
          "execution_count": 7,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Check the lengths\n",
        "len(train_data), len(test_data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "f1UxnvvmblUj",
        "outputId": "1dd0c011-912e-498b-ea34-ee4bf05c7b87"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(<torch.utils.data.dataloader.DataLoader at 0x7fca2e344760>,\n",
              " <torch.utils.data.dataloader.DataLoader at 0x7fca2e3445b0>)"
            ]
          },
          "execution_count": 8,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Turn train and test Datasets into DataLoaders\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "train_dataloader = DataLoader(dataset=train_data, \n",
        "                              batch_size=1, # how many samples per batch?\n",
        "                              num_workers=1, # how many subprocesses to use for data loading? (higher = more)\n",
        "                              shuffle=True) # shuffle the data?\n",
        "\n",
        "test_dataloader = DataLoader(dataset=test_data, \n",
        "                             batch_size=1, \n",
        "                             num_workers=1, \n",
        "                             shuffle=False) # don't usually need to shuffle testing data\n",
        "\n",
        "train_dataloader, test_dataloader"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JJKRqXNGbnI5",
        "outputId": "2d5320fc-ac72-4ef8-94f8-0ef283fc3a30"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Image shape: torch.Size([1, 3, 64, 64]) -> [batch_size, color_channels, height, width]\n",
            "Label shape: torch.Size([1])\n"
          ]
        }
      ],
      "source": [
        "# Check out single image size/shape\n",
        "img, label = next(iter(train_dataloader))\n",
        "\n",
        "# Batch size will now be 1, try changing the batch_size parameter above and see what happens\n",
        "print(f\"Image shape: {img.shape} -> [batch_size, color_channels, height, width]\")\n",
        "print(f\"Label shape: {label.shape}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RjINjcyGdwL9"
      },
      "source": [
        "### 2.1 Create Datasets and DataLoaders (script mode)\n",
        "\n",
        "Rather than rewriting all of the code above everytime we wanted to load data, we can turn it into a script called `data_setup.py`.\n",
        "\n",
        "Let's capture all of the above functionality into a function called `create_dataloaders()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CmjK-b04du_-",
        "outputId": "07d3450d-bc24-4aa4-a6d6-57d60bf11416"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Overwriting going_modular/data_setup.py\n"
          ]
        }
      ],
      "source": [
        "%%writefile going_modular/data_setup.py\n",
        "\"\"\"\n",
        "Contains functionality for creating PyTorch DataLoaders for \n",
        "image classification data.\n",
        "\"\"\"\n",
        "import os\n",
        "\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision import datasets, transforms\n",
        "\n",
        "NUM_WORKERS = os.cpu_count()\n",
        "\n",
        "def create_dataloaders(\n",
        "    train_dir: str, \n",
        "    test_dir: str, \n",
        "    transform: transforms.Compose, \n",
        "    batch_size: int, \n",
        "    num_workers: int=NUM_WORKERS\n",
        "):\n",
        "  \"\"\"Creates training and testing DataLoaders.\n",
        "\n",
        "  Takes in a training directory and testing directory path and turns\n",
        "  them into PyTorch Datasets and then into PyTorch DataLoaders.\n",
        "\n",
        "  Args:\n",
        "    train_dir: Path to training directory.\n",
        "    test_dir: Path to testing directory.\n",
        "    transform: torchvision transforms to perform on training and testing data.\n",
        "    batch_size: Number of samples per batch in each of the DataLoaders.\n",
        "    num_workers: An integer for number of workers per DataLoader.\n",
        "\n",
        "  Returns:\n",
        "    A tuple of (train_dataloader, test_dataloader, class_names).\n",
        "    Where class_names is a list of the target classes.\n",
        "    Example usage:\n",
        "      train_dataloader, test_dataloader, class_names = \\\n",
        "        = create_dataloaders(train_dir=path/to/train_dir,\n",
        "                             test_dir=path/to/test_dir,\n",
        "                             transform=some_transform,\n",
        "                             batch_size=32,\n",
        "                             num_workers=4)\n",
        "  \"\"\"\n",
        "  # Use ImageFolder to create dataset(s)\n",
        "  train_data = datasets.ImageFolder(train_dir, transform=transform)\n",
        "  test_data = datasets.ImageFolder(test_dir, transform=transform)\n",
        "\n",
        "  # Get class names\n",
        "  class_names = train_data.classes\n",
        "\n",
        "  # Turn images into data loaders\n",
        "  train_dataloader = DataLoader(\n",
        "      train_data,\n",
        "      batch_size=batch_size,\n",
        "      shuffle=True,\n",
        "      num_workers=num_workers,\n",
        "      pin_memory=True,\n",
        "  )\n",
        "  test_dataloader = DataLoader(\n",
        "      test_data,\n",
        "      batch_size=batch_size,\n",
        "      shuffle=False,\n",
        "      num_workers=num_workers,\n",
        "      pin_memory=True,\n",
        "  )\n",
        "\n",
        "  return train_dataloader, test_dataloader, class_names"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gAQLp1dCbrY0"
      },
      "source": [
        "## 3. Making a model (TinyVGG) \n",
        "\n",
        "We're going to use the same model we used in notebook 04: TinyVGG from the CNN Explainer website.\n",
        "\n",
        "The only change here from notebook 04 is that a docstring has been added using [Google's Style Guide for Python](https://google.github.io/styleguide/pyguide.html#384-classes). "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z6i2xWcwbvom"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "\n",
        "from torch import nn \n",
        "\n",
        "class TinyVGG(nn.Module):\n",
        "    \"\"\"Creates the TinyVGG architecture.\n",
        "\n",
        "    Replicates the TinyVGG architecture from the CNN explainer website in PyTorch.\n",
        "    See the original architecture here: https://poloclub.github.io/cnn-explainer/\n",
        "\n",
        "    Args:\n",
        "    input_shape: An integer indicating number of input channels.\n",
        "    hidden_units: An integer indicating number of hidden units between layers.\n",
        "    output_shape: An integer indicating number of output units.\n",
        "    \"\"\"\n",
        "    def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:\n",
        "        super().__init__()\n",
        "        self.conv_block_1 = nn.Sequential(\n",
        "          nn.Conv2d(in_channels=input_shape, \n",
        "                    out_channels=hidden_units, \n",
        "                    kernel_size=3, \n",
        "                    stride=1, \n",
        "                    padding=0),  \n",
        "          nn.ReLU(),\n",
        "          nn.Conv2d(in_channels=hidden_units, \n",
        "                    out_channels=hidden_units,\n",
        "                    kernel_size=3,\n",
        "                    stride=1,\n",
        "                    padding=0),\n",
        "          nn.ReLU(),\n",
        "          nn.MaxPool2d(kernel_size=2,\n",
        "                        stride=2)\n",
        "        )\n",
        "        self.conv_block_2 = nn.Sequential(\n",
        "          nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),\n",
        "          nn.ReLU(),\n",
        "          nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),\n",
        "          nn.ReLU(),\n",
        "          nn.MaxPool2d(2)\n",
        "        )\n",
        "        self.classifier = nn.Sequential(\n",
        "          nn.Flatten(),\n",
        "          # Where did this in_features shape come from? \n",
        "          # It's because each layer of our network compresses and changes the shape of our inputs data.\n",
        "          nn.Linear(in_features=hidden_units*13*13,\n",
        "                    out_features=output_shape)\n",
        "        )\n",
        "    \n",
        "    def forward(self, x: torch.Tensor):\n",
        "        x = self.conv_block_1(x)\n",
        "        x = self.conv_block_2(x)\n",
        "        x = self.classifier(x)\n",
        "        return x\n",
        "        # return self.classifier(self.block_2(self.block_1(x))) # <- leverage the benefits of operator fusion"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ROZsju1d9Cbu"
      },
      "source": [
        "Now let's create an instance of `TinyVGG` and put it on the target device.\n",
        "\n",
        "> **Note:** If you're using Google Colab, and you'd like to use a GPU (recommended), you can turn one on via going to Runtime -> Change runtime type -> Hardware accelerator -> GPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tXOOAEHHc6-L",
        "outputId": "3207fda6-896b-4cb5-9f08-f1e46879bf75"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "TinyVGG(\n",
              "  (conv_block_1): Sequential(\n",
              "    (0): Conv2d(3, 10, kernel_size=(3, 3), stride=(1, 1))\n",
              "    (1): ReLU()\n",
              "    (2): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1))\n",
              "    (3): ReLU()\n",
              "    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
              "  )\n",
              "  (conv_block_2): Sequential(\n",
              "    (0): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1))\n",
              "    (1): ReLU()\n",
              "    (2): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1))\n",
              "    (3): ReLU()\n",
              "    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
              "  )\n",
              "  (classifier): Sequential(\n",
              "    (0): Flatten(start_dim=1, end_dim=-1)\n",
              "    (1): Linear(in_features=1690, out_features=3, bias=True)\n",
              "  )\n",
              ")"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import torch\n",
        "\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n",
        "# Instantiate an instance of the model\n",
        "torch.manual_seed(42)\n",
        "model_0 = TinyVGG(input_shape=3, # number of color channels (3 for RGB) \n",
        "                  hidden_units=10, \n",
        "                  output_shape=len(train_data.classes)).to(device)\n",
        "model_0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C1CwvPiMcEfZ"
      },
      "source": [
        "Let's check out our model by doing a dummy forward pass."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "__p5G2ArcFko",
        "outputId": "341f2b73-a0fa-4fdb-f458-142510dcfcca"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Single image shape: torch.Size([1, 3, 64, 64])\n",
            "\n",
            "Output logits:\n",
            "tensor([[ 0.0208, -0.0019,  0.0095]], device='cuda:0')\n",
            "\n",
            "Output prediction probabilities:\n",
            "tensor([[0.3371, 0.3295, 0.3333]], device='cuda:0')\n",
            "\n",
            "Output prediction label:\n",
            "tensor([0], device='cuda:0')\n",
            "\n",
            "Actual label:\n",
            "0\n"
          ]
        }
      ],
      "source": [
        "# 1. Get a batch of images and labels from the DataLoader\n",
        "img_batch, label_batch = next(iter(train_dataloader))\n",
        "\n",
        "# 2. Get a single image from the batch and unsqueeze the image so its shape fits the model\n",
        "img_single, label_single = img_batch[0].unsqueeze(dim=0), label_batch[0]\n",
        "print(f\"Single image shape: {img_single.shape}\\n\")\n",
        "\n",
        "# 3. Perform a forward pass on a single image\n",
        "model_0.eval()\n",
        "with torch.inference_mode():\n",
        "    pred = model_0(img_single.to(device))\n",
        "    \n",
        "# 4. Print out what's happening and convert model logits -> pred probs -> pred label\n",
        "print(f\"Output logits:\\n{pred}\\n\")\n",
        "print(f\"Output prediction probabilities:\\n{torch.softmax(pred, dim=1)}\\n\")\n",
        "print(f\"Output prediction label:\\n{torch.argmax(torch.softmax(pred, dim=1), dim=1)}\\n\")\n",
        "print(f\"Actual label:\\n{label_single}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wSaDm_W7bc3y"
      },
      "source": [
        "### 3.1 Making a model (TinyVGG) (script mode)\n",
        "\n",
        "Over the past few notebooks (notebook 03 and notebook 04), we've built the TinyVGG model a few times.\n",
        "\n",
        "So it makes sense to put the model into its file so we can reuse it again and again.\n",
        "\n",
        "Let's put our `TinyVGG()` model class into a script called `model_builder.py` with the line `%%writefile going_modular/model_builder.py`. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oEH3bklL86hF",
        "outputId": "b3f7e983-e51c-4e35-a991-df5ff99d63e6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Overwriting going_modular/model_builder.py\n"
          ]
        }
      ],
      "source": [
        "%%writefile going_modular/model_builder.py\n",
        "\"\"\"\n",
        "Contains PyTorch model code to instantiate a TinyVGG model.\n",
        "\"\"\"\n",
        "import torch\n",
        "\n",
        "from torch import nn\n",
        "\n",
        "class TinyVGG(nn.Module):\n",
        "    \"\"\"Creates the TinyVGG architecture.\n",
        "\n",
        "    Replicates the TinyVGG architecture from the CNN explainer website in PyTorch.\n",
        "    See the original architecture here: https://poloclub.github.io/cnn-explainer/\n",
        "\n",
        "    Args:\n",
        "    input_shape: An integer indicating number of input channels.\n",
        "    hidden_units: An integer indicating number of hidden units between layers.\n",
        "    output_shape: An integer indicating number of output units.\n",
        "    \"\"\"\n",
        "    def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:\n",
        "        super().__init__()\n",
        "        self.conv_block_1 = nn.Sequential(\n",
        "          nn.Conv2d(in_channels=input_shape, \n",
        "                    out_channels=hidden_units, \n",
        "                    kernel_size=3, \n",
        "                    stride=1, \n",
        "                    padding=0),  \n",
        "          nn.ReLU(),\n",
        "          nn.Conv2d(in_channels=hidden_units, \n",
        "                    out_channels=hidden_units,\n",
        "                    kernel_size=3,\n",
        "                    stride=1,\n",
        "                    padding=0),\n",
        "          nn.ReLU(),\n",
        "          nn.MaxPool2d(kernel_size=2,\n",
        "                        stride=2)\n",
        "        )\n",
        "        self.conv_block_2 = nn.Sequential(\n",
        "          nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),\n",
        "          nn.ReLU(),\n",
        "          nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),\n",
        "          nn.ReLU(),\n",
        "          nn.MaxPool2d(2)\n",
        "        )\n",
        "        self.classifier = nn.Sequential(\n",
        "          nn.Flatten(),\n",
        "          # Where did this in_features shape come from? \n",
        "          # It's because each layer of our network compresses and changes the shape of our inputs data.\n",
        "          nn.Linear(in_features=hidden_units*13*13,\n",
        "                    out_features=output_shape)\n",
        "        )\n",
        "    \n",
        "    def forward(self, x: torch.Tensor):\n",
        "        x = self.conv_block_1(x)\n",
        "        x = self.conv_block_2(x)\n",
        "        x = self.classifier(x)\n",
        "        return x\n",
        "        # return self.classifier(self.block_2(self.block_1(x))) # <- leverage the benefits of operator fusion"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3-WijBWM9K2c"
      },
      "source": [
        "Create an instance of `TinyVGG` (from the script)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "F__hbQwZ9A56",
        "outputId": "4e8abd04-cbc4-4aa8-b71c-3a8ff8617407"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "TinyVGG(\n",
              "  (conv_block_1): Sequential(\n",
              "    (0): Conv2d(3, 10, kernel_size=(3, 3), stride=(1, 1))\n",
              "    (1): ReLU()\n",
              "    (2): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1))\n",
              "    (3): ReLU()\n",
              "    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
              "  )\n",
              "  (conv_block_2): Sequential(\n",
              "    (0): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1))\n",
              "    (1): ReLU()\n",
              "    (2): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1))\n",
              "    (3): ReLU()\n",
              "    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
              "  )\n",
              "  (classifier): Sequential(\n",
              "    (0): Flatten(start_dim=1, end_dim=-1)\n",
              "    (1): Linear(in_features=1690, out_features=3, bias=True)\n",
              "  )\n",
              ")"
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import torch\n",
        "\n",
        "from going_modular import model_builder\n",
        "\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n",
        "# Instantiate an instance of the model from the \"model_builder.py\" script\n",
        "torch.manual_seed(42)\n",
        "model_1 = model_builder.TinyVGG(input_shape=3, # number of color channels (3 for RGB) \n",
        "                                hidden_units=10, \n",
        "                                output_shape=len(class_names)).to(device)\n",
        "model_1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vBD-yWcp9QVT"
      },
      "source": [
        "Do a dummy forward pass on `model_1`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NCQWeeMj9kHE",
        "outputId": "cb33c4ec-4848-423c-aea3-89a131e07cfd"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Single image shape: torch.Size([1, 3, 64, 64])\n",
            "\n",
            "Output logits:\n",
            "tensor([[ 0.0208, -0.0019,  0.0095]], device='cuda:0')\n",
            "\n",
            "Output prediction probabilities:\n",
            "tensor([[0.3371, 0.3295, 0.3333]], device='cuda:0')\n",
            "\n",
            "Output prediction label:\n",
            "tensor([0], device='cuda:0')\n",
            "\n",
            "Actual label:\n",
            "0\n"
          ]
        }
      ],
      "source": [
        "# 1. Get a batch of images and labels from the DataLoader\n",
        "img_batch, label_batch = next(iter(train_dataloader))\n",
        "\n",
        "# 2. Get a single image from the batch and unsqueeze the image so its shape fits the model\n",
        "img_single, label_single = img_batch[0].unsqueeze(dim=0), label_batch[0]\n",
        "print(f\"Single image shape: {img_single.shape}\\n\")\n",
        "\n",
        "# 3. Perform a forward pass on a single image\n",
        "model_1.eval()\n",
        "with torch.inference_mode():\n",
        "    pred = model_1(img_single.to(device))\n",
        "    \n",
        "# 4. Print out what's happening and convert model logits -> pred probs -> pred label\n",
        "print(f\"Output logits:\\n{pred}\\n\")\n",
        "print(f\"Output prediction probabilities:\\n{torch.softmax(pred, dim=1)}\\n\")\n",
        "print(f\"Output prediction label:\\n{torch.argmax(torch.softmax(pred, dim=1), dim=1)}\\n\")\n",
        "print(f\"Actual label:\\n{label_single}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8OB4xIsXcHrp"
      },
      "source": [
        "## 4. Creating `train_step()` and `test_step()` functions and `train()` to combine them  \n",
        "\n",
        "Rather than writing them again, we can reuse the `train_step()` and `test_step()` functions from [notebook 04](https://www.learnpytorch.io/04_pytorch_custom_datasets/#75-create-train-test-loop-functions).\n",
        "\n",
        "The same goes for the `train()` function we created.\n",
        "\n",
        "The only difference here is that these functions have had docstrings added to them in [Google's Python Functions and Methods Style Guide](https://google.github.io/styleguide/pyguide.html#383-functions-and-methods).\n",
        "\n",
        "Let's start by making `train_step()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t_0zRYZecJZ6"
      },
      "outputs": [],
      "source": [
        "from typing import Tuple\n",
        "\n",
        "def train_step(model: torch.nn.Module, \n",
        "               dataloader: torch.utils.data.DataLoader, \n",
        "               loss_fn: torch.nn.Module, \n",
        "               optimizer: torch.optim.Optimizer,\n",
        "               device: torch.device) -> Tuple[float, float]:\n",
        "    \"\"\"Trains a PyTorch model for a single epoch.\n",
        "\n",
        "    Turns a target PyTorch model to training mode and then\n",
        "    runs through all of the required training steps (forward\n",
        "    pass, loss calculation, optimizer step).\n",
        "\n",
        "    Args:\n",
        "    model: A PyTorch model to be trained.\n",
        "    dataloader: A DataLoader instance for the model to be trained on.\n",
        "    loss_fn: A PyTorch loss function to minimize.\n",
        "    optimizer: A PyTorch optimizer to help minimize the loss function.\n",
        "    device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
        "\n",
        "    Returns:\n",
        "    A tuple of training loss and training accuracy metrics.\n",
        "    In the form (train_loss, train_accuracy). For example:\n",
        "\n",
        "    (0.1112, 0.8743)\n",
        "    \"\"\"\n",
        "    # Put model in train mode\n",
        "    model.train()\n",
        "\n",
        "    # Setup train loss and train accuracy values\n",
        "    train_loss, train_acc = 0, 0\n",
        "\n",
        "    # Loop through data loader data batches\n",
        "    for batch, (X, y) in enumerate(dataloader):\n",
        "        # Send data to target device\n",
        "        X, y = X.to(device), y.to(device)\n",
        "\n",
        "        # 1. Forward pass\n",
        "        y_pred = model(X)\n",
        "\n",
        "        # 2. Calculate  and accumulate loss\n",
        "        loss = loss_fn(y_pred, y)\n",
        "        train_loss += loss.item() \n",
        "\n",
        "        # 3. Optimizer zero grad\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        # 4. Loss backward\n",
        "        loss.backward()\n",
        "\n",
        "        # 5. Optimizer step\n",
        "        optimizer.step()\n",
        "\n",
        "        # Calculate and accumulate accuracy metric across all batches\n",
        "        y_pred_class = torch.argmax(torch.softmax(y_pred, dim=1), dim=1)\n",
        "        train_acc += (y_pred_class == y).sum().item()/len(y_pred)\n",
        "\n",
        "    # Adjust metrics to get average loss and accuracy per batch \n",
        "    train_loss = train_loss / len(dataloader)\n",
        "    train_acc = train_acc / len(dataloader)\n",
        "    return train_loss, train_acc"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1nwAPwNh-f3k"
      },
      "source": [
        "Now we'll do `test_step()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sr_9AspYcKkW"
      },
      "outputs": [],
      "source": [
        "def test_step(model: torch.nn.Module, \n",
        "              dataloader: torch.utils.data.DataLoader, \n",
        "              loss_fn: torch.nn.Module,\n",
        "              device: torch.device) -> Tuple[float, float]:\n",
        "    \"\"\"Tests a PyTorch model for a single epoch.\n",
        "\n",
        "    Turns a target PyTorch model to \"eval\" mode and then performs\n",
        "    a forward pass on a testing dataset.\n",
        "\n",
        "    Args:\n",
        "    model: A PyTorch model to be tested.\n",
        "    dataloader: A DataLoader instance for the model to be tested on.\n",
        "    loss_fn: A PyTorch loss function to calculate loss on the test data.\n",
        "    device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
        "\n",
        "    Returns:\n",
        "    A tuple of testing loss and testing accuracy metrics.\n",
        "    In the form (test_loss, test_accuracy). For example:\n",
        "\n",
        "    (0.0223, 0.8985)\n",
        "    \"\"\"\n",
        "    # Put model in eval mode\n",
        "    model.eval() \n",
        "\n",
        "    # Setup test loss and test accuracy values\n",
        "    test_loss, test_acc = 0, 0\n",
        "\n",
        "    # Turn on inference context manager\n",
        "    with torch.inference_mode():\n",
        "        # Loop through DataLoader batches\n",
        "        for batch, (X, y) in enumerate(dataloader):\n",
        "            # Send data to target device\n",
        "            X, y = X.to(device), y.to(device)\n",
        "\n",
        "            # 1. Forward pass\n",
        "            test_pred_logits = model(X)\n",
        "\n",
        "            # 2. Calculate and accumulate loss\n",
        "            loss = loss_fn(test_pred_logits, y)\n",
        "            test_loss += loss.item()\n",
        "\n",
        "            # Calculate and accumulate accuracy\n",
        "            test_pred_labels = test_pred_logits.argmax(dim=1)\n",
        "            test_acc += ((test_pred_labels == y).sum().item()/len(test_pred_labels))\n",
        "\n",
        "    # Adjust metrics to get average loss and accuracy per batch \n",
        "    test_loss = test_loss / len(dataloader)\n",
        "    test_acc = test_acc / len(dataloader)\n",
        "    return test_loss, test_acc"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "weqHDd5b-iGz"
      },
      "source": [
        "And we'll combine `train_step()` and `test_step()` into `train()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-oze8b6icWE7"
      },
      "outputs": [],
      "source": [
        "from typing import Dict, List\n",
        "\n",
        "from tqdm.auto import tqdm\n",
        "\n",
        "def train(model: torch.nn.Module, \n",
        "          train_dataloader: torch.utils.data.DataLoader, \n",
        "          test_dataloader: torch.utils.data.DataLoader, \n",
        "          optimizer: torch.optim.Optimizer,\n",
        "          loss_fn: torch.nn.Module,\n",
        "          epochs: int,\n",
        "          device: torch.device) -> Dict[str, List[float]]:\n",
        "    \"\"\"Trains and tests a PyTorch model.\n",
        "\n",
        "    Passes a target PyTorch models through train_step() and test_step()\n",
        "    functions for a number of epochs, training and testing the model\n",
        "    in the same epoch loop.\n",
        "\n",
        "    Calculates, prints and stores evaluation metrics throughout.\n",
        "\n",
        "    Args:\n",
        "    model: A PyTorch model to be trained and tested.\n",
        "    train_dataloader: A DataLoader instance for the model to be trained on.\n",
        "    test_dataloader: A DataLoader instance for the model to be tested on.\n",
        "    optimizer: A PyTorch optimizer to help minimize the loss function.\n",
        "    loss_fn: A PyTorch loss function to calculate loss on both datasets.\n",
        "    epochs: An integer indicating how many epochs to train for.\n",
        "    device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
        "\n",
        "    Returns:\n",
        "    A dictionary of training and testing loss as well as training and\n",
        "    testing accuracy metrics. Each metric has a value in a list for \n",
        "    each epoch.\n",
        "    In the form: {train_loss: [...],\n",
        "                  train_acc: [...],\n",
        "                  test_loss: [...],\n",
        "                  test_acc: [...]} \n",
        "    For example if training for epochs=2: \n",
        "                 {train_loss: [2.0616, 1.0537],\n",
        "                  train_acc: [0.3945, 0.3945],\n",
        "                  test_loss: [1.2641, 1.5706],\n",
        "                  test_acc: [0.3400, 0.2973]} \n",
        "    \"\"\"\n",
        "    # Create empty results dictionary\n",
        "    results = {\"train_loss\": [],\n",
        "      \"train_acc\": [],\n",
        "      \"test_loss\": [],\n",
        "      \"test_acc\": []\n",
        "    }\n",
        "\n",
        "    # Loop through training and testing steps for a number of epochs\n",
        "    for epoch in tqdm(range(epochs)):\n",
        "        train_loss, train_acc = train_step(model=model,\n",
        "                                          dataloader=train_dataloader,\n",
        "                                          loss_fn=loss_fn,\n",
        "                                          optimizer=optimizer,\n",
        "                                          device=device)\n",
        "        test_loss, test_acc = test_step(model=model,\n",
        "          dataloader=test_dataloader,\n",
        "          loss_fn=loss_fn,\n",
        "          device=device)\n",
        "\n",
        "        # Print out what's happening\n",
        "        print(\n",
        "          f\"Epoch: {epoch+1} | \"\n",
        "          f\"train_loss: {train_loss:.4f} | \"\n",
        "          f\"train_acc: {train_acc:.4f} | \"\n",
        "          f\"test_loss: {test_loss:.4f} | \"\n",
        "          f\"test_acc: {test_acc:.4f}\"\n",
        "        )\n",
        "\n",
        "        # Update results dictionary\n",
        "        results[\"train_loss\"].append(train_loss)\n",
        "        results[\"train_acc\"].append(train_acc)\n",
        "        results[\"test_loss\"].append(test_loss)\n",
        "        results[\"test_acc\"].append(test_acc)\n",
        "\n",
        "    # Return the filled results at the end of the epochs\n",
        "    return results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tw_nYa4kfQeI"
      },
      "source": [
        "### 4.1 Creating `train_step()` and `test_step()` functions and `train()` to combine them (script mode)   \n",
        "\n",
        "To create a script for `train_step()`, `test_step()` and `train()`, we'll combine their code all into a single cell.\n",
        "\n",
        "We'll then write that cell to a file called `engine.py` because these functions will be the \"engine\" of our training pipeline.\n",
        "\n",
        "We can do so with the magic line `%%writefile going_modular/engine.py`.\n",
        "\n",
        "We'll also make sure to put all the imports we need (`torch`, `typing`, and `tqdm`) at the top of the cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "d47bhfY3fVq5",
        "outputId": "41053838-2e28-4702-d833-d8c97eed119b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Overwriting going_modular/engine.py\n"
          ]
        }
      ],
      "source": [
        "%%writefile going_modular/engine.py\n",
        "\"\"\"\n",
        "Contains functions for training and testing a PyTorch model.\n",
        "\"\"\"\n",
        "from typing import Dict, List, Tuple\n",
        "\n",
        "import torch\n",
        "\n",
        "from tqdm.auto import tqdm\n",
        "\n",
        "def train_step(model: torch.nn.Module, \n",
        "               dataloader: torch.utils.data.DataLoader, \n",
        "               loss_fn: torch.nn.Module, \n",
        "               optimizer: torch.optim.Optimizer,\n",
        "               device: torch.device) -> Tuple[float, float]:\n",
        "    \"\"\"Trains a PyTorch model for a single epoch.\n",
        "\n",
        "    Turns a target PyTorch model to training mode and then\n",
        "    runs through all of the required training steps (forward\n",
        "    pass, loss calculation, optimizer step).\n",
        "\n",
        "    Args:\n",
        "    model: A PyTorch model to be trained.\n",
        "    dataloader: A DataLoader instance for the model to be trained on.\n",
        "    loss_fn: A PyTorch loss function to minimize.\n",
        "    optimizer: A PyTorch optimizer to help minimize the loss function.\n",
        "    device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
        "\n",
        "    Returns:\n",
        "    A tuple of training loss and training accuracy metrics.\n",
        "    In the form (train_loss, train_accuracy). For example:\n",
        "\n",
        "    (0.1112, 0.8743)\n",
        "    \"\"\"\n",
        "    # Put model in train mode\n",
        "    model.train()\n",
        "\n",
        "    # Setup train loss and train accuracy values\n",
        "    train_loss, train_acc = 0, 0\n",
        "\n",
        "    # Loop through data loader data batches\n",
        "    for batch, (X, y) in enumerate(dataloader):\n",
        "        # Send data to target device\n",
        "        X, y = X.to(device), y.to(device)\n",
        "\n",
        "        # 1. Forward pass\n",
        "        y_pred = model(X)\n",
        "\n",
        "        # 2. Calculate  and accumulate loss\n",
        "        loss = loss_fn(y_pred, y)\n",
        "        train_loss += loss.item() \n",
        "\n",
        "        # 3. Optimizer zero grad\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        # 4. Loss backward\n",
        "        loss.backward()\n",
        "\n",
        "        # 5. Optimizer step\n",
        "        optimizer.step()\n",
        "\n",
        "        # Calculate and accumulate accuracy metric across all batches\n",
        "        y_pred_class = torch.argmax(torch.softmax(y_pred, dim=1), dim=1)\n",
        "        train_acc += (y_pred_class == y).sum().item()/len(y_pred)\n",
        "\n",
        "    # Adjust metrics to get average loss and accuracy per batch \n",
        "    train_loss = train_loss / len(dataloader)\n",
        "    train_acc = train_acc / len(dataloader)\n",
        "    return train_loss, train_acc\n",
        "\n",
        "def test_step(model: torch.nn.Module, \n",
        "              dataloader: torch.utils.data.DataLoader, \n",
        "              loss_fn: torch.nn.Module,\n",
        "              device: torch.device) -> Tuple[float, float]:\n",
        "    \"\"\"Tests a PyTorch model for a single epoch.\n",
        "\n",
        "    Turns a target PyTorch model to \"eval\" mode and then performs\n",
        "    a forward pass on a testing dataset.\n",
        "\n",
        "    Args:\n",
        "    model: A PyTorch model to be tested.\n",
        "    dataloader: A DataLoader instance for the model to be tested on.\n",
        "    loss_fn: A PyTorch loss function to calculate loss on the test data.\n",
        "    device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
        "\n",
        "    Returns:\n",
        "    A tuple of testing loss and testing accuracy metrics.\n",
        "    In the form (test_loss, test_accuracy). For example:\n",
        "\n",
        "    (0.0223, 0.8985)\n",
        "    \"\"\"\n",
        "    # Put model in eval mode\n",
        "    model.eval() \n",
        "\n",
        "    # Setup test loss and test accuracy values\n",
        "    test_loss, test_acc = 0, 0\n",
        "\n",
        "    # Turn on inference context manager\n",
        "    with torch.inference_mode():\n",
        "        # Loop through DataLoader batches\n",
        "        for batch, (X, y) in enumerate(dataloader):\n",
        "            # Send data to target device\n",
        "            X, y = X.to(device), y.to(device)\n",
        "\n",
        "            # 1. Forward pass\n",
        "            test_pred_logits = model(X)\n",
        "\n",
        "            # 2. Calculate and accumulate loss\n",
        "            loss = loss_fn(test_pred_logits, y)\n",
        "            test_loss += loss.item()\n",
        "\n",
        "            # Calculate and accumulate accuracy\n",
        "            test_pred_labels = test_pred_logits.argmax(dim=1)\n",
        "            test_acc += ((test_pred_labels == y).sum().item()/len(test_pred_labels))\n",
        "\n",
        "    # Adjust metrics to get average loss and accuracy per batch \n",
        "    test_loss = test_loss / len(dataloader)\n",
        "    test_acc = test_acc / len(dataloader)\n",
        "    return test_loss, test_acc\n",
        "\n",
        "def train(model: torch.nn.Module, \n",
        "          train_dataloader: torch.utils.data.DataLoader, \n",
        "          test_dataloader: torch.utils.data.DataLoader, \n",
        "          optimizer: torch.optim.Optimizer,\n",
        "          loss_fn: torch.nn.Module,\n",
        "          epochs: int,\n",
        "          device: torch.device) -> Dict[str, List[float]]:\n",
        "    \"\"\"Trains and tests a PyTorch model.\n",
        "\n",
        "    Passes a target PyTorch models through train_step() and test_step()\n",
        "    functions for a number of epochs, training and testing the model\n",
        "    in the same epoch loop.\n",
        "\n",
        "    Calculates, prints and stores evaluation metrics throughout.\n",
        "\n",
        "    Args:\n",
        "    model: A PyTorch model to be trained and tested.\n",
        "    train_dataloader: A DataLoader instance for the model to be trained on.\n",
        "    test_dataloader: A DataLoader instance for the model to be tested on.\n",
        "    optimizer: A PyTorch optimizer to help minimize the loss function.\n",
        "    loss_fn: A PyTorch loss function to calculate loss on both datasets.\n",
        "    epochs: An integer indicating how many epochs to train for.\n",
        "    device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
        "\n",
        "    Returns:\n",
        "    A dictionary of training and testing loss as well as training and\n",
        "    testing accuracy metrics. Each metric has a value in a list for \n",
        "    each epoch.\n",
        "    In the form: {train_loss: [...],\n",
        "              train_acc: [...],\n",
        "              test_loss: [...],\n",
        "              test_acc: [...]} \n",
        "    For example if training for epochs=2: \n",
        "             {train_loss: [2.0616, 1.0537],\n",
        "              train_acc: [0.3945, 0.3945],\n",
        "              test_loss: [1.2641, 1.5706],\n",
        "              test_acc: [0.3400, 0.2973]} \n",
        "    \"\"\"\n",
        "    # Create empty results dictionary\n",
        "    results = {\"train_loss\": [],\n",
        "               \"train_acc\": [],\n",
        "               \"test_loss\": [],\n",
        "               \"test_acc\": []\n",
        "    }\n",
        "\n",
        "    # Loop through training and testing steps for a number of epochs\n",
        "    for epoch in tqdm(range(epochs)):\n",
        "        train_loss, train_acc = train_step(model=model,\n",
        "                                          dataloader=train_dataloader,\n",
        "                                          loss_fn=loss_fn,\n",
        "                                          optimizer=optimizer,\n",
        "                                          device=device)\n",
        "        test_loss, test_acc = test_step(model=model,\n",
        "          dataloader=test_dataloader,\n",
        "          loss_fn=loss_fn,\n",
        "          device=device)\n",
        "\n",
        "        # Print out what's happening\n",
        "        print(\n",
        "          f\"Epoch: {epoch+1} | \"\n",
        "          f\"train_loss: {train_loss:.4f} | \"\n",
        "          f\"train_acc: {train_acc:.4f} | \"\n",
        "          f\"test_loss: {test_loss:.4f} | \"\n",
        "          f\"test_acc: {test_acc:.4f}\"\n",
        "        )\n",
        "\n",
        "        # Update results dictionary\n",
        "        results[\"train_loss\"].append(train_loss)\n",
        "        results[\"train_acc\"].append(train_acc)\n",
        "        results[\"test_loss\"].append(test_loss)\n",
        "        results[\"test_acc\"].append(test_acc)\n",
        "\n",
        "    # Return the filled results at the end of the epochs\n",
        "    return results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A2nXumITc9DL"
      },
      "source": [
        "## 5. Creating a function to save the model\n",
        "\n",
        "Let's setup a function to save our model to a directory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lu_vyjq-c933"
      },
      "outputs": [],
      "source": [
        "from pathlib import Path\n",
        "\n",
        "def save_model(model: torch.nn.Module,\n",
        "               target_dir: str,\n",
        "               model_name: str):\n",
        "    \"\"\"Saves a PyTorch model to a target directory.\n",
        "\n",
        "    Args:\n",
        "    model: A target PyTorch model to save.\n",
        "    target_dir: A directory for saving the model to.\n",
        "    model_name: A filename for the saved model. Should include\n",
        "      either \".pth\" or \".pt\" as the file extension.\n",
        "\n",
        "    Example usage:\n",
        "    save_model(model=model_0,\n",
        "               target_dir=\"models\",\n",
        "               model_name=\"05_going_modular_tingvgg_model.pth\")\n",
        "    \"\"\"\n",
        "    # Create target directory\n",
        "    target_dir_path = Path(target_dir)\n",
        "    target_dir_path.mkdir(parents=True,\n",
        "                        exist_ok=True)\n",
        "\n",
        "    # Create model save path\n",
        "    assert model_name.endswith(\".pth\") or model_name.endswith(\".pt\"), \"model_name should end with '.pt' or '.pth'\"\n",
        "    model_save_path = target_dir_path / model_name\n",
        "\n",
        "    # Save the model state_dict()\n",
        "    print(f\"[INFO] Saving model to: {model_save_path}\")\n",
        "    torch.save(obj=model.state_dict(),\n",
        "             f=model_save_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F7SnPGOfc_Wc"
      },
      "source": [
        "### 5.1 Creating a function to save the model (script mode)\n",
        "\n",
        "How about we add our `save_model()` function to a script called `utils.py` which is short for \"utilities\".\n",
        "\n",
        "We can do so with the magic line `%%writefile going_modular/utils.py`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "20LC4xwN_2ZT",
        "outputId": "30852196-8138-4f9e-9f43-578dc6277682"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Overwriting going_modular/utils.py\n"
          ]
        }
      ],
      "source": [
        "%%writefile going_modular/utils.py\n",
        "\"\"\"\n",
        "Contains various utility functions for PyTorch model training and saving.\n",
        "\"\"\"\n",
        "from pathlib import Path\n",
        "\n",
        "import torch\n",
        "\n",
        "def save_model(model: torch.nn.Module,\n",
        "               target_dir: str,\n",
        "               model_name: str):\n",
        "    \"\"\"Saves a PyTorch model to a target directory.\n",
        "\n",
        "    Args:\n",
        "    model: A target PyTorch model to save.\n",
        "    target_dir: A directory for saving the model to.\n",
        "    model_name: A filename for the saved model. Should include\n",
        "      either \".pth\" or \".pt\" as the file extension.\n",
        "\n",
        "    Example usage:\n",
        "    save_model(model=model_0,\n",
        "               target_dir=\"models\",\n",
        "               model_name=\"05_going_modular_tingvgg_model.pth\")\n",
        "    \"\"\"\n",
        "    # Create target directory\n",
        "    target_dir_path = Path(target_dir)\n",
        "    target_dir_path.mkdir(parents=True,\n",
        "                        exist_ok=True)\n",
        "\n",
        "    # Create model save path\n",
        "    assert model_name.endswith(\".pth\") or model_name.endswith(\".pt\"), \"model_name should end with '.pt' or '.pth'\"\n",
        "    model_save_path = target_dir_path / model_name\n",
        "\n",
        "    # Save the model state_dict()\n",
        "    print(f\"[INFO] Saving model to: {model_save_path}\")\n",
        "    torch.save(obj=model.state_dict(),\n",
        "             f=model_save_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "20ZO64U2cZY_"
      },
      "source": [
        "## 6. Train, evaluate and save the model\n",
        "\n",
        "Let's leverage the functions we've got above to train, test and save a model to file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 173,
          "referenced_widgets": [
            "84317b054c244ee4b1bed7b2f86660b7",
            "a0a72aed7e2743f8a95c6af3b6e4a457",
            "b517f855551f4c53969677161952a1ac",
            "4d780e75e5d24241b1f6c0203a8bd73e",
            "a0caf8f4760c4935b5b9be86566ee8e0",
            "83cf71b79abd4fc5ae9daaac8f586841",
            "3402f79d3df643efbed25c72aac94642",
            "1e79615e722f456189e07874d13564d2",
            "bac27d87b4d54b6cbaf85dd7b7ecdc64",
            "f9b7fbae53e84883ac281343bb6a9da9",
            "4bc100911d7d4303b6ab7efb59d3420c",
            "205f478a42ca42f2ab32113d597dc388"
          ]
        },
        "id": "FjeRMiLjccVc",
        "outputId": "2d9005a4-afce-4230-c758-1c6efa01a5e8"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "205f478a42ca42f2ab32113d597dc388",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "  0%|          | 0/5 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch: 1 | train_loss: 1.0956 | train_acc: 0.3867 | test_loss: 1.0630 | test_acc: 0.4133\n",
            "Epoch: 2 | train_loss: 1.0141 | train_acc: 0.5111 | test_loss: 1.0210 | test_acc: 0.4533\n",
            "Epoch: 3 | train_loss: 0.9591 | train_acc: 0.5644 | test_loss: 0.9961 | test_acc: 0.4400\n",
            "Epoch: 4 | train_loss: 0.8994 | train_acc: 0.5778 | test_loss: 0.9986 | test_acc: 0.4533\n",
            "Epoch: 5 | train_loss: 0.8652 | train_acc: 0.6267 | test_loss: 1.0010 | test_acc: 0.5467\n",
            "[INFO] Total training time: 5.461 seconds\n",
            "[INFO] Saving model to: models/05_going_modular_cell_mode_tinyvgg_model.pth\n"
          ]
        }
      ],
      "source": [
        "# Set random seeds\n",
        "torch.manual_seed(42) \n",
        "torch.cuda.manual_seed(42)\n",
        "\n",
        "# Set number of epochs\n",
        "NUM_EPOCHS = 5\n",
        "\n",
        "# Recreate an instance of TinyVGG\n",
        "model_0 = TinyVGG(input_shape=3, # number of color channels (3 for RGB) \n",
        "                  hidden_units=10, \n",
        "                  output_shape=len(train_data.classes)).to(device)\n",
        "\n",
        "# Setup loss function and optimizer\n",
        "loss_fn = nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.Adam(params=model_0.parameters(), lr=0.001)\n",
        "\n",
        "# Start the timer\n",
        "from timeit import default_timer as timer \n",
        "start_time = timer()\n",
        "\n",
        "# Train model_0 \n",
        "model_0_results = train(model=model_0, \n",
        "                        train_dataloader=train_dataloader,\n",
        "                        test_dataloader=test_dataloader,\n",
        "                        optimizer=optimizer,\n",
        "                        loss_fn=loss_fn, \n",
        "                        epochs=NUM_EPOCHS,\n",
        "                        device=device)\n",
        "\n",
        "# End the timer and print out how long it took\n",
        "end_time = timer()\n",
        "print(f\"[INFO] Total training time: {end_time-start_time:.3f} seconds\")\n",
        "\n",
        "# Save the model\n",
        "save_model(model=model_0,\n",
        "           target_dir=\"models\",\n",
        "           model_name=\"05_going_modular_cell_mode_tinyvgg_model.pth\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Md_CmXTvfkOz"
      },
      "source": [
        "### 6.1 Train, evaluate and save the model (script mode)\n",
        "\n",
        "Let's combine all of our modular files into a single script `train.py`.\n",
        "\n",
        "This will allow us to run all of the functions we've written with a single line of code on the command line:\n",
        "\n",
        "`python going_modular/train.py`\n",
        "\n",
        "Or if we're running it in a notebook:\n",
        "\n",
        "`!python going_modular/train.py`\n",
        "\n",
        "We'll go through the following steps:\n",
        "1. Import the various dependencies, namely `torch`, `os`, `torchvision.transforms` and all of the scripts from the `going_modular` directory, `data_setup`, `engine`, `model_builder`, `utils`.\n",
        "  * **Note:** Since `train.py` will be *inside* the `going_modular` directory, we can import the other modules via `import ...` rather than `from going_modular import ...`.\n",
        "2. Setup various hyperparameters such as batch size, number of epochs, learning rate and number of hidden units (these could be set in the future via [Python's `argparse`](https://docs.python.org/3/library/argparse.html)).\n",
        "3. Setup the training and test directories.\n",
        "4. Setup device-agnostic code.\n",
        "5. Create the necessary data transforms.\n",
        "6. Create the DataLoaders using `data_setup.py`.\n",
        "7. Create the model using `model_builder.py`.\n",
        "8. Setup the loss function and optimizer.\n",
        "9. Train the model using `engine.py`.\n",
        "10. Save the model using `utils.py`. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8XnG2l4zf2Ei",
        "outputId": "84f67ba0-ddaf-45a2-9a71-a37cf3d1ca88"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Overwriting going_modular/train.py\n"
          ]
        }
      ],
      "source": [
        "%%writefile going_modular/train.py\n",
        "\"\"\"\n",
        "Trains a PyTorch image classification model using device-agnostic code.\n",
        "\"\"\"\n",
        "\n",
        "import os\n",
        "\n",
        "import torch\n",
        "\n",
        "from torchvision import transforms\n",
        "\n",
        "import data_setup, engine, model_builder, utils\n",
        "\n",
        "\n",
        "# Setup hyperparameters\n",
        "NUM_EPOCHS = 5\n",
        "BATCH_SIZE = 32\n",
        "HIDDEN_UNITS = 10\n",
        "LEARNING_RATE = 0.001\n",
        "\n",
        "# Setup directories\n",
        "train_dir = \"data/pizza_steak_sushi/train\"\n",
        "test_dir = \"data/pizza_steak_sushi/test\"\n",
        "\n",
        "# Setup target device\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n",
        "# Create transforms\n",
        "data_transform = transforms.Compose([\n",
        "  transforms.Resize((64, 64)),\n",
        "  transforms.ToTensor()\n",
        "])\n",
        "\n",
        "# Create DataLoaders with help from data_setup.py\n",
        "train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(\n",
        "    train_dir=train_dir,\n",
        "    test_dir=test_dir,\n",
        "    transform=data_transform,\n",
        "    batch_size=BATCH_SIZE\n",
        ")\n",
        "\n",
        "# Create model with help from model_builder.py\n",
        "model = model_builder.TinyVGG(\n",
        "    input_shape=3,\n",
        "    hidden_units=HIDDEN_UNITS,\n",
        "    output_shape=len(class_names)\n",
        ").to(device)\n",
        "\n",
        "# Set loss and optimizer\n",
        "loss_fn = torch.nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.Adam(model.parameters(),\n",
        "                             lr=LEARNING_RATE)\n",
        "\n",
        "# Start training with help from engine.py\n",
        "engine.train(model=model,\n",
        "             train_dataloader=train_dataloader,\n",
        "             test_dataloader=test_dataloader,\n",
        "             loss_fn=loss_fn,\n",
        "             optimizer=optimizer,\n",
        "             epochs=NUM_EPOCHS,\n",
        "             device=device)\n",
        "\n",
        "# Save the model with help from utils.py\n",
        "utils.save_model(model=model,\n",
        "                 target_dir=\"models\",\n",
        "                 model_name=\"05_going_modular_script_mode_tinyvgg_model.pth\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HHnPm-w7CjXT"
      },
      "source": [
        "Now our final directory structure looks like:\n",
        "```\n",
        "data/\n",
        "  pizza_steak_sushi/\n",
        "    train/\n",
        "      pizza/\n",
        "        train_image_01.jpeg\n",
        "        train_image_02.jpeg\n",
        "        ...\n",
        "      steak/\n",
        "      sushi/\n",
        "    test/\n",
        "      pizza/\n",
        "        test_image_01.jpeg\n",
        "        test_image_02.jpeg\n",
        "        ...\n",
        "      steak/\n",
        "      sushi/\n",
        "going_modular/\n",
        "  data_setup.py\n",
        "  engine.py\n",
        "  model_builder.py\n",
        "  train.py\n",
        "  utils.py\n",
        "models/\n",
        "  saved_model.pth\n",
        "```\n",
        "\n",
        "Now to put it all together!\n",
        "\n",
        "Let's run our `train.py` file from the command line with:\n",
        "\n",
        "```\n",
        "!python going_modular/train.py\n",
        "```\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6eVtanbSCjFj",
        "outputId": "9c90d0d8-9b44-4558-b74d-292729eca1e4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "  0%|                                                     | 0/5 [00:00<?, ?it/s]Epoch: 1 | train_loss: 1.1131 | train_acc: 0.2852 | test_loss: 1.1138 | test_acc: 0.2604\n",
            " 20%|█████████                                    | 1/5 [00:01<00:04,  1.06s/it]Epoch: 2 | train_loss: 1.0851 | train_acc: 0.4102 | test_loss: 1.1238 | test_acc: 0.1979\n",
            " 40%|██████████████████                           | 2/5 [00:01<00:02,  1.20it/s]Epoch: 3 | train_loss: 1.0837 | train_acc: 0.4141 | test_loss: 1.1459 | test_acc: 0.1979\n",
            " 60%|███████████████████████████                  | 3/5 [00:02<00:01,  1.33it/s]Epoch: 4 | train_loss: 1.1104 | train_acc: 0.2930 | test_loss: 1.1318 | test_acc: 0.1979\n",
            " 80%|████████████████████████████████████         | 4/5 [00:03<00:00,  1.40it/s]Epoch: 5 | train_loss: 1.0833 | train_acc: 0.2930 | test_loss: 1.0883 | test_acc: 0.3712\n",
            "100%|█████████████████████████████████████████████| 5/5 [00:03<00:00,  1.35it/s]\n",
            "[INFO] Saving model to: models/05_going_modular_script_mode_tinyvgg_model.pth\n"
          ]
        }
      ],
      "source": [
        "!python going_modular/train.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vJdtyrk0V5_e"
      },
      "source": [
        "\n",
        "We've just trained a model with a single line of code from the command line.\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "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"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "1e79615e722f456189e07874d13564d2": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "3402f79d3df643efbed25c72aac94642": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "4bc100911d7d4303b6ab7efb59d3420c": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "4d780e75e5d24241b1f6c0203a8bd73e": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_f9b7fbae53e84883ac281343bb6a9da9",
            "placeholder": "​",
            "style": "IPY_MODEL_4bc100911d7d4303b6ab7efb59d3420c",
            "value": " 5/5 [00:11&lt;00:00,  2.23s/it]"
          }
        },
        "83cf71b79abd4fc5ae9daaac8f586841": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "84317b054c244ee4b1bed7b2f86660b7": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_a0a72aed7e2743f8a95c6af3b6e4a457",
              "IPY_MODEL_b517f855551f4c53969677161952a1ac",
              "IPY_MODEL_4d780e75e5d24241b1f6c0203a8bd73e"
            ],
            "layout": "IPY_MODEL_a0caf8f4760c4935b5b9be86566ee8e0"
          }
        },
        "a0a72aed7e2743f8a95c6af3b6e4a457": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_83cf71b79abd4fc5ae9daaac8f586841",
            "placeholder": "​",
            "style": "IPY_MODEL_3402f79d3df643efbed25c72aac94642",
            "value": "100%"
          }
        },
        "a0caf8f4760c4935b5b9be86566ee8e0": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "b517f855551f4c53969677161952a1ac": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_1e79615e722f456189e07874d13564d2",
            "max": 5,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_bac27d87b4d54b6cbaf85dd7b7ecdc64",
            "value": 5
          }
        },
        "bac27d87b4d54b6cbaf85dd7b7ecdc64": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "f9b7fbae53e84883ac281343bb6a9da9": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}