11_PyTorch_Going_Modular.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# <center><b><font color=\"red\">PyTorch Going Modular</font></b></center>"
      ],
      "metadata": {
        "id": "DaMJLsmkFlQO"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Script Mode**"
      ],
      "metadata": {
        "id": "pWdSCL-EOdSS"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "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",
        "> 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."
      ],
      "metadata": {
        "id": "7AOdPN4Rd-VD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import torch"
      ],
      "metadata": {
        "id": "amLt25S9t0bW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vnD9aZENE8vt",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "30d14bbe-7094-4cc9-ef9b-a9886f99794b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing hello_world.py\n"
          ]
        }
      ],
      "source": [
        "# Create a Python file called hello_world.py\n",
        "%%writefile hello_world.py\n",
        "print(\"I Love PyTorch\")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Run this Python file\n",
        "!python hello_world.py"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LJfs9Yq0OTja",
        "outputId": "0f60fed0-fe25-4de9-d92b-2a892a9fc213"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "I Love PyTorch\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "! rm -rf hello_world.py"
      ],
      "metadata": {
        "id": "HhZwf1J6OE7L"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Directory Structure**\n",
        "\n",
        "⭐ **PyTorch project structure:**\n",
        "```\n",
        "pytorch_project/\n",
        "├── pytorch_project/\n",
        "│   ├── data_setup.py\n",
        "│   ├── engine.py\n",
        "│   ├── model.py\n",
        "│   ├── train.py\n",
        "│   └── utils.py\n",
        "├── models/\n",
        "│   ├── model_1.pth\n",
        "│   └── model_2.pth\n",
        "└── data/\n",
        "    ├── data_folder_1/\n",
        "    └── data_folder_2/\n",
        "```\n",
        "\n",
        "⭐ `TinyVGG` ~ `pizza_steak_sushi`\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",
        "    "
      ],
      "metadata": {
        "id": "1tF-d54Ledy_"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **1. Get data**"
      ],
      "metadata": {
        "id": "jKgQxoCwlBzq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "os.makedirs(\"going_modular\", exist_ok=True)"
      ],
      "metadata": {
        "id": "wVhu1ko4fy2A"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import requests\n",
        "import zipfile\n",
        "from pathlib import Path\n",
        "\n",
        "# Step 1: Set up the path\n",
        "url = \"https://github.com/mohd-faizy/PyTorch-Essentials/raw/main/_datasets/pizza_steak_sushi.zip\"\n",
        "image_path = Path(\"data/pizza_steak_sushi\")\n",
        "image_path.mkdir(parents=True, exist_ok=True)\n",
        "zip_path = image_path.parent / \"pizza_steak_sushi.zip\"\n",
        "\n",
        "# Step 2: Download the zip file containing the data\n",
        "zip_path.write_bytes(requests.get(url).content)\n",
        "\n",
        "# # Step 3: Unzip the downloaded file into the image directory\n",
        "with zipfile.ZipFile(zip_path, \"r\") as zip_ref:\n",
        "    zip_ref.extractall(image_path)"
      ],
      "metadata": {
        "id": "AECZFkSclEbf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!ls /content/data"
      ],
      "metadata": {
        "id": "MHLUHhc_lNmt",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2f0f8968-5e14-4469-cf15-2772ee111cb5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "pizza_steak_sushi  pizza_steak_sushi.zip\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# setup train and testing paths\n",
        "train_dir = image_path / \"train\"\n",
        "test_dir = image_path / \"test\"\n",
        "\n",
        "train_dir, test_dir"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "E8moIESYinTs",
        "outputId": "5270e301-b50f-4ef1-b28e-0c4f99d51a74"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(PosixPath('data/pizza_steak_sushi/train'),\n",
              " PosixPath('data/pizza_steak_sushi/test'))"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **2. Create Datasets and DataLoaders**"
      ],
      "metadata": {
        "id": "qV9wAAGWlE3I"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "`transforms`  → `ImageFolder(datasets)` → `DataLoader`"
      ],
      "metadata": {
        "id": "gFNANT78URC-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%writefile going_modular/data_setup.py\n",
        "\"\"\"\n",
        "Provides functionality for creating PyTorch DataLoaders for\n",
        "image classification tasks.\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",
        "    \"\"\"Generates training and testing DataLoaders from directory paths.\n",
        "\n",
        "    Converts the specified training and testing directories into\n",
        "    PyTorch Datasets and then into DataLoaders.\n",
        "\n",
        "    Args:\n",
        "        train_dir (str): Path to the training data directory.\n",
        "        test_dir (str): Path to the testing data directory.\n",
        "        transform (transforms.Compose): Transformations to apply to the data.\n",
        "        batch_size (int): Number of samples per batch in each DataLoader.\n",
        "        num_workers (int, optional): Number of subprocesses to use for data loading.\n",
        "                                     Defaults to the number of CPU cores.\n",
        "\n",
        "    Returns:\n",
        "        tuple: Contains (train_dataloader, test_dataloader, class_names).\n",
        "            - train_dataloader (DataLoader): DataLoader for the training data.\n",
        "            - test_dataloader (DataLoader): DataLoader for the testing data.\n",
        "            - class_names (list): List of class names from the target dataset.\n",
        "\n",
        "    Example:\n",
        "        train_dataloader, test_dataloader, class_names = create_dataloaders(\n",
        "            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",
        "    \"\"\"\n",
        "\n",
        "    # Create datasets from image folders\n",
        "    train_data = datasets.ImageFolder(train_dir,\n",
        "                                      transform=transform)\n",
        "    test_data = datasets.ImageFolder(test_dir,\n",
        "                                     transform=transform)\n",
        "\n",
        "    # Extract class names\n",
        "    class_names = train_data.classes\n",
        "\n",
        "    # Initialize DataLoaders\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"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xJ_TKtnEW_NA",
        "outputId": "19d5e462-4d21-49e3-a28f-448c01798d91"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing going_modular/data_setup.py\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from going_modular import data_setup\n",
        "from torchvision import transforms\n",
        "\n",
        "# Create simple transform\n",
        "data_transform = transforms.Compose([\n",
        "    transforms.Resize((64, 64)),\n",
        "    transforms.ToTensor()\n",
        "])\n",
        "\n",
        "NUM_WORKERS = os.cpu_count()\n",
        "\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=1,\n",
        "    num_workers=os.cpu_count()\n",
        ")\n",
        "\n",
        "image, label = next(iter(train_dataloader))\n",
        "image.shape, label.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FxgXU8ubcojt",
        "outputId": "3bbc7fe7-1a82-430c-d424-835daafd95b7"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(torch.Size([1, 3, 64, 64]), torch.Size([1]))"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **3. Making a model (TinyVGG)**\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "DRlNtUaBmYIe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%writefile going_modular/model_builder.py\n",
        "\"\"\"\n",
        "Provides PyTorch model code for instantiating a TinyVGG model.\n",
        "\"\"\"\n",
        "import torch\n",
        "from torch import nn\n",
        "\n",
        "class TinyVGG(nn.Module):\n",
        "    \"\"\"Defines the TinyVGG architecture.\n",
        "\n",
        "    This class replicates the TinyVGG architecture as described on the CNN explainer website.\n",
        "    The original architecture can be found here: https://poloclub.github.io/cnn-explainer/\n",
        "\n",
        "    Args:\n",
        "        input_shape (int): Number of input channels.\n",
        "        hidden_units (int): Number of hidden units between layers.\n",
        "        output_shape (int): 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, 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",
        "            nn.Linear(in_features=hidden_units * 13 * 13, out_features=output_shape)\n",
        "        )\n",
        "\n",
        "    def forward(self, x: torch.Tensor) -> 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.conv_block_2(self.conv_block_1(x))) # <- leverage the benefits of operator fusion"
      ],
      "metadata": {
        "id": "ncPqjGJ1mZnu",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "bc338f69-6a1b-4078-8240-61b6efc063b1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing going_modular/model_builder.py\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "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 = 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"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aKF_VlH9cgns",
        "outputId": "5b0fb05b-58d5-4f03-a0cd-712311010726"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "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",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Do a dummy forward pass on model.**"
      ],
      "metadata": {
        "id": "EE9vQcrIgeAc"
      }
    },
    {
      "cell_type": "code",
      "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.eval()\n",
        "with torch.inference_mode():\n",
        "    pred = model(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}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GlaxxNLUffIw",
        "outputId": "6713f7c9-5c96-4771-bb09-9fdb0456573c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Single image shape: torch.Size([1, 3, 64, 64])\n",
            "\n",
            "Output logits:\n",
            "tensor([[ 0.0208, -0.0020,  0.0095]])\n",
            "\n",
            "Output prediction probabilities:\n",
            "tensor([[0.3371, 0.3295, 0.3333]])\n",
            "\n",
            "Output prediction label:\n",
            "tensor([0])\n",
            "\n",
            "Actual label:\n",
            "0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **4. Creating `train_step()` and `test_step()` functions and `train()` to combine them.**"
      ],
      "metadata": {
        "id": "AWvHhJIYmZ5B"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%writefile going_modular/engine.py\n",
        "\"\"\"\n",
        "Contains functions for training and testing a PyTorch model.\n",
        "\"\"\"\n",
        "\n",
        "from typing import Dict, List, Tuple\n",
        "import torch\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",
        "    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 in the form\n",
        "        (train_loss, train_accuracy).\n",
        "    \"\"\"\n",
        "    # Set model to training mode\n",
        "    model.train()\n",
        "\n",
        "    # Initialize training loss and accuracy\n",
        "    train_loss, train_acc = 0, 0\n",
        "\n",
        "    # Loop through data batches\n",
        "    for batch, (X, y) in enumerate(dataloader):\n",
        "        # Transfer data to target device\n",
        "        X, y = X.to(device), y.to(device)\n",
        "\n",
        "        # Forward pass\n",
        "        y_pred = model(X)\n",
        "\n",
        "        # Calculate and accumulate loss\n",
        "        loss = loss_fn(y_pred, y)\n",
        "        train_loss += loss.item()\n",
        "\n",
        "        # Zero gradients\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        # Backward pass\n",
        "        loss.backward()\n",
        "\n",
        "        # Optimizer step\n",
        "        optimizer.step()\n",
        "\n",
        "        # Calculate and accumulate accuracy\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",
        "    # Calculate average loss and accuracy per batch\n",
        "    train_loss /= len(dataloader)\n",
        "    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",
        "    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 in the form\n",
        "        (test_loss, test_accuracy).\n",
        "    \"\"\"\n",
        "    # Set model to evaluation mode\n",
        "    model.eval()\n",
        "\n",
        "    # Initialize test loss and accuracy\n",
        "    test_loss, test_acc = 0, 0\n",
        "\n",
        "    # Disable gradient calculations for inference\n",
        "    with torch.inference_mode():\n",
        "        # Loop through data batches\n",
        "        for batch, (X, y) in enumerate(dataloader):\n",
        "            # Transfer data to target device\n",
        "            X, y = X.to(device), y.to(device)\n",
        "\n",
        "            # Forward pass\n",
        "            test_pred_logits = model(X)\n",
        "\n",
        "            # 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",
        "    # Calculate average loss and accuracy per batch\n",
        "    test_loss /= len(dataloader)\n",
        "    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 over multiple epochs.\n",
        "\n",
        "    Args:\n",
        "        model: A PyTorch model to be trained and tested.\n",
        "        train_dataloader: A DataLoader instance for training data.\n",
        "        test_dataloader: A DataLoader instance for testing data.\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: Number of 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 metrics, each containing a list\n",
        "        of values for each epoch. Format:\n",
        "        {\n",
        "            \"train_loss\": [...],\n",
        "            \"train_acc\": [...],\n",
        "            \"test_loss\": [...],\n",
        "            \"test_acc\": [...]\n",
        "        }\n",
        "    \"\"\"\n",
        "    # Initialize results dictionary\n",
        "    results = {\n",
        "        \"train_loss\": [],\n",
        "        \"train_acc\": [],\n",
        "        \"test_loss\": [],\n",
        "        \"test_acc\": []\n",
        "    }\n",
        "\n",
        "    # Loop through epochs\n",
        "    for epoch in tqdm(range(epochs)):\n",
        "        # Training step\n",
        "        train_loss, train_acc = train_step(\n",
        "            model=model,\n",
        "            dataloader=train_dataloader,\n",
        "            loss_fn=loss_fn,\n",
        "            optimizer=optimizer,\n",
        "            device=device\n",
        "        )\n",
        "\n",
        "        # Testing step\n",
        "        test_loss, test_acc = test_step(\n",
        "            model=model,\n",
        "            dataloader=test_dataloader,\n",
        "            loss_fn=loss_fn,\n",
        "            device=device\n",
        "        )\n",
        "\n",
        "        # Print progress\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 results\n",
        "    return results"
      ],
      "metadata": {
        "id": "M_FIPuu8mlqe",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "79e2e015-0f43-4157-f5ad-589e1d76faf7"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing going_modular/engine.py\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **5. Creating a function to save the model**"
      ],
      "metadata": {
        "id": "uJYIOLI3mmCX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%writefile going_modular/utils.py\n",
        "\"\"\"\n",
        "Contains various utility functions for PyTorch model training and saving.\n",
        "\"\"\"\n",
        "import os\n",
        "import torch\n",
        "from pathlib import Path\n",
        "\n",
        "def save_model(model: torch.nn.Module,\n",
        "               target_dir: str,\n",
        "               model_name: str):\n",
        "    \"\"\"\n",
        "    Saves a PyTorch model to a specified directory.\n",
        "\n",
        "    Args:\n",
        "        model (torch.nn.Module): The PyTorch model to be saved.\n",
        "        target_dir (str): The directory where the model will be saved.\n",
        "        model_name (str): The name of the saved model file. Must include \".pth\" or \".pt\" as the file extension.\n",
        "\n",
        "    Example:\n",
        "        save_model(model=model_0,\n",
        "        target_dir=\"models\",\n",
        "        model_name=\"05_going_modular_tingvgg_model.pth\")\n",
        "    \"\"\"\n",
        "    # Create the target directory if it doesn't exist\n",
        "    target_dir_path = Path(target_dir)\n",
        "    target_dir_path.mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "    # Ensure the model name has the correct file extension\n",
        "    assert model_name.endswith(\".pth\") or model_name.endswith(\".pt\"), \"model_name must end with '.pt' or '.pth'\"\n",
        "\n",
        "    # Define the full path for the saved model\n",
        "    model_save_path = target_dir_path / model_name\n",
        "\n",
        "    # Save the model's state dictionary\n",
        "    print(f\"[INFO] Saving model to: {model_save_path}\")\n",
        "    torch.save(obj=model.state_dict(),\n",
        "               f=model_save_path)"
      ],
      "metadata": {
        "id": "TIdye3UznqKG",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2458dac7-f38a-4831-cdbd-fa5ddde0d3e6"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing going_modular/utils.py\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **6. Train, evaluate and save the model**"
      ],
      "metadata": {
        "id": "AVYHNXbinqiu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%writefile going_modular/train.py\n",
        "\"\"\"\n",
        "Trains a PyTorch image classification model using device-agnostic code.\n",
        "\"\"\"\n",
        "\n",
        "import os\n",
        "import torch\n",
        "from torchvision import transforms\n",
        "import data_setup, engine, model_builder, utils\n",
        "\n",
        "\n",
        "# Setup hyperparameters\n",
        "NUM_EPOCHS = 10\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",
        "# 1. 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",
        "# 2. 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",
        "# 3. 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",
        "# 4. 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\")"
      ],
      "metadata": {
        "id": "SYjRx5IFntKe",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "549d8dcd-93f3-4c6f-ddd1-bdf4346ec3ce"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing going_modular/train.py\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "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",
        "```"
      ],
      "metadata": {
        "id": "y2l-ER0stbtg"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "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",
        "```"
      ],
      "metadata": {
        "id": "wt1LKsXXtn3k"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# 5. Train----------------------------------------------------------------------\n",
        "!python going_modular/train.py"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dZQEXKyctjkd",
        "outputId": "9476fa0e-3c59-4e3a-df3a-71f5685778e9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "  0% 0/10 [00:00<?, ?it/s]Epoch: 1 | train_loss: 1.1029 | train_acc: 0.2891 | test_loss: 1.0995 | test_acc: 0.2604\n",
            " 10% 1/10 [00:02<00:26,  2.91s/it]Epoch: 2 | train_loss: 1.0940 | train_acc: 0.4375 | test_loss: 1.0818 | test_acc: 0.5417\n",
            " 20% 2/10 [00:04<00:18,  2.34s/it]Epoch: 3 | train_loss: 1.1077 | train_acc: 0.2812 | test_loss: 1.0711 | test_acc: 0.5417\n",
            " 30% 3/10 [00:06<00:14,  2.11s/it]Epoch: 4 | train_loss: 1.1007 | train_acc: 0.2812 | test_loss: 1.0778 | test_acc: 0.5521\n",
            " 40% 4/10 [00:08<00:12,  2.01s/it]Epoch: 5 | train_loss: 1.0838 | train_acc: 0.4844 | test_loss: 1.0801 | test_acc: 0.4848\n",
            " 50% 5/10 [00:10<00:09,  1.93s/it]Epoch: 6 | train_loss: 1.0770 | train_acc: 0.4727 | test_loss: 1.0700 | test_acc: 0.4328\n",
            " 60% 6/10 [00:12<00:07,  1.89s/it]Epoch: 7 | train_loss: 1.0534 | train_acc: 0.4805 | test_loss: 1.0553 | test_acc: 0.3419\n",
            " 70% 7/10 [00:14<00:06,  2.15s/it]Epoch: 8 | train_loss: 1.0071 | train_acc: 0.4961 | test_loss: 1.0454 | test_acc: 0.3930\n",
            " 80% 8/10 [00:17<00:04,  2.20s/it]Epoch: 9 | train_loss: 0.9537 | train_acc: 0.5977 | test_loss: 0.9851 | test_acc: 0.4536\n",
            " 90% 9/10 [00:19<00:02,  2.10s/it]Epoch: 10 | train_loss: 0.9213 | train_acc: 0.5000 | test_loss: 0.9891 | test_acc: 0.3920\n",
            "100% 10/10 [00:20<00:00,  2.09s/it]\n",
            "[INFO] Saving model to: models/05_going_modular_script_mode_tinyvgg_model.pth\n"
          ]
        }
      ]
    }
  ]
}