05_pytorch_going_modular.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyO0cikPiwsrLFomMYYR6/E6",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/patrickabadi/python-pytorch-notebooks/blob/main/05_pytorch_going_modular.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "os.makedirs(\"going_modular\", exist_ok=True)"
      ],
      "metadata": {
        "id": "0xJch2rrPDXe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Q63-zYS4NCfr",
        "outputId": "db31b97f-4b39-47f6-f67c-87f6f89b5edd"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Did not find data/pizza_steak_sushi directory, creating one...\n",
            "Downloading pizza, steak, sushi data...\n",
            "Unzipping pizza, steak, sushi data...\n"
          ]
        }
      ],
      "source": [
        "import os\n",
        "import requests\n",
        "import zipfile\n",
        "from pathlib import Path\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",
      "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": "Bgvx_rkPPhC9",
        "outputId": "a3d85613-427a-4717-ee74-bd01a332996b"
      },
      "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": 5
        }
      ]
    },
    {
      "cell_type": "code",
      "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",
        "])"
      ],
      "metadata": {
        "id": "m6GK_to0QHJ9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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 torchvision import datasets, transforms\n",
        "from torch.utils.data import DataLoader\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=True,\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": "uCFkKraxNtJ9",
        "outputId": "ce99cfee-84ec-40d2-b049-93474e9fcac5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing going_modular/data_setup.py\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from going_modular import data_setup\n",
        "\n",
        "train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(train_dir = train_dir,\n",
        "                                                                               test_dir = test_dir,\n",
        "                                                                               transform=data_transform,\n",
        "                                                                               batch_size = 32,\n",
        "                                                                               num_workers=2)"
      ],
      "metadata": {
        "id": "UbJJbsOQOLRs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%writefile going_modular/model_builder.py\n",
        "\"\"\"\n",
        "Contains PyTorch model code to instantiate a TinyVGG model.\n",
        "\"\"\"\n",
        "import torch\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.conv_block_2(self.conv_block_1(x))) # <- leverage the benefits of operator fusion\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ffJmOSK1Oh69",
        "outputId": "efd19494-3e8a-48e3-9656-856c6b23be16"
      },
      "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",
        "torch.manual_seed(42)\n",
        "model = model_builder.TinyVGG(input_shape=3,\n",
        "                              hidden_units = 10,\n",
        "                              output_shape=len(class_names)).to(device)"
      ],
      "metadata": {
        "id": "OumUdxSLO7m6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%writefile going_modular/engine.py\n",
        "\"\"\"\n",
        "Contains functions for training and testing a PyTorch model.\n",
        "\"\"\"\n",
        "import torch\n",
        "\n",
        "from tqdm.auto import tqdm\n",
        "from typing import Dict, List, 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\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]:\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"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TgL4d4nxSHpN",
        "outputId": "7fdeea75-f46d-46dc-8c9d-a8527fa3835f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing going_modular/engine.py\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#from going_modular import engine\n",
        "\n",
        "#engine.train(model=model,\n",
        " #            train_dataloader=train_dataloader,\n",
        "  #           test_dataloader=test_dataloader,)"
      ],
      "metadata": {
        "id": "pIzpFFqBUY9u"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SH9h61b1U6dV",
        "outputId": "2a29744e-4f97-453e-c1e7-5c7e41812e6e"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing going_modular/utils.py\n"
          ]
        }
      ]
    },
    {
      "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",
        "import data_setup, engine, model_builder, utils\n",
        "\n",
        "from torchvision import transforms\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\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XlV99tKdV0D7",
        "outputId": "1be2054b-3b9a-4733-b26a-54ee9f5671b1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing going_modular/train.py\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!python going_modular/train.py"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YccFYg58WCvz",
        "outputId": "309c34b2-dd72-4382-ffef-32a4b51adb5b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "  0% 0/5 [00:00<?, ?it/s]Epoch: 1 | train_loss: 1.1034 | train_acc: 0.3945 | test_loss: 1.0920 | test_acc: 0.3797\n",
            " 20% 1/5 [00:08<00:32,  8.19s/it]Epoch: 2 | train_loss: 1.1145 | train_acc: 0.3047 | test_loss: 1.0904 | test_acc: 0.3400\n",
            " 40% 2/5 [00:10<00:13,  4.53s/it]Epoch: 3 | train_loss: 1.0921 | train_acc: 0.3242 | test_loss: 1.0895 | test_acc: 0.4934\n",
            " 60% 3/5 [00:12<00:06,  3.39s/it]Epoch: 4 | train_loss: 1.0954 | train_acc: 0.3281 | test_loss: 1.0868 | test_acc: 0.2680\n",
            " 80% 4/5 [00:13<00:02,  2.64s/it]Epoch: 5 | train_loss: 1.0183 | train_acc: 0.5195 | test_loss: 1.0372 | test_acc: 0.4839\n",
            "100% 5/5 [00:15<00:00,  3.03s/it]\n",
            "[INFO] Saving model to: models/05_going_modular_script_mode_tinyvgg_model.pth\n"
          ]
        }
      ]
    }
  ]
}