Going Modular

from pathlib import Path

Path("going_modular").mkdir(exist_ok=True)

def beam_area(width: float, depth: float) -> float:
    """Returns the cross-sectional area of a rectangular beam.

    Args:
        width: Width of the beam in meters.
        depth: Depth of the beam in meters.

    Returns:
        The area in square meters.
    """
    return width * depth

help(beam_area)

%%writefile going_modular/get_data.py
"""
Downloads the pizza, steak and sushi images into data/pizza_steak_sushi/.
"""
import io
import zipfile
from pathlib import Path

import requests

URL = ("https://raw.githubusercontent.com/mrdbourke/pytorch-deep-learning/"
       "main/data/pizza_steak_sushi.zip")

image_path = Path("data") / "pizza_steak_sushi"

if image_path.is_dir():
    print(f"{image_path} already exists, skipping download.")
else:
    print(f"Downloading {URL}")
    response = requests.get(URL)
    response.raise_for_status()   # stop with an error if the download failed

    # Unzip straight from memory into data/pizza_steak_sushi/
    print(f"Unzipping to {image_path}")
    with zipfile.ZipFile(io.BytesIO(response.content)) as zip_file:
        zip_file.extractall(image_path)
    print("Done.")

!python going_modular/get_data.py

for split in ["train", "test"]:
    for class_dir in sorted((Path("data/pizza_steak_sushi") / split).iterdir()):
        print(f"{split}/{class_dir.name}: {len(list(class_dir.glob('*.jpg')))} images")

%%writefile going_modular/data_setup.py
"""
Functions for turning folders of images into PyTorch DataLoaders
for image classification.
"""
import os

import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

NUM_WORKERS = os.cpu_count()


def create_dataloaders(
    train_dir: str,
    test_dir: str,
    transform: transforms.Compose,
    batch_size: int,
    num_workers: int = NUM_WORKERS,
):
    """Creates training and testing DataLoaders.

    Reads images from a training folder and a testing folder, each laid out
    in the standard image-classification format (one subfolder per class),
    and turns them into PyTorch Datasets and then DataLoaders.

    Args:
        train_dir: Path to the training folder.
        test_dir: Path to the testing folder.
        transform: torchvision transforms to apply to every image.
        batch_size: Number of images per batch in each DataLoader.
        num_workers: Number of worker processes that load batches in parallel.

    Returns:
        A tuple (train_dataloader, test_dataloader, class_names), where
        class_names is a list of the target classes.

    Example:
        train_dataloader, test_dataloader, class_names = create_dataloaders(
            train_dir="data/pizza_steak_sushi/train",
            test_dir="data/pizza_steak_sushi/test",
            transform=some_transform,
            batch_size=32,
        )
    """
    # ImageFolder turns "one subfolder per class" into a labeled dataset
    train_data = datasets.ImageFolder(train_dir, transform=transform)
    test_data = datasets.ImageFolder(test_dir, transform=transform)

    # Class names come from the subfolder names, e.g. ["pizza", "steak", "sushi"]
    class_names = train_data.classes

    # Pinned memory speeds up copying batches to a GPU; it is only useful when there is one
    pin_memory = torch.cuda.is_available()

    train_dataloader = DataLoader(
        train_data,
        batch_size=batch_size,
        shuffle=True,             # shuffle the training data every epoch
        num_workers=num_workers,
        pin_memory=pin_memory,
    )
    test_dataloader = DataLoader(
        test_data,
        batch_size=batch_size,
        shuffle=False,            # keep the test data in a fixed order
        num_workers=num_workers,
        pin_memory=pin_memory,
    )

    return train_dataloader, test_dataloader, class_names

%%writefile going_modular/model_builder.py
"""
PyTorch model code for a TinyVGG convolutional neural network.
"""
import torch
from torch import nn


class TinyVGG(nn.Module):
    """Creates the TinyVGG architecture.

    Replicates the TinyVGG architecture from the CNN Explainer website:
    https://poloclub.github.io/cnn-explainer/

    Args:
        input_shape: Number of input color channels (3 for RGB images).
        hidden_units: Number of filters in each convolutional layer.
        output_shape: Number of output classes.
    """

    def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:
        super().__init__()
        self.conv_block_1 = nn.Sequential(
            nn.Conv2d(in_channels=input_shape, out_channels=hidden_units,
                      kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units,
                      kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        self.conv_block_2 = nn.Sequential(
            nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),
            nn.ReLU(),
            nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            # A 64x64 input is 13x13 after the two blocks (see module 04)
            nn.Linear(in_features=hidden_units * 13 * 13, out_features=output_shape),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv_block_1(x)
        x = self.conv_block_2(x)
        x = self.classifier(x)
        return x

%%writefile going_modular/engine.py
"""
Functions for training and testing a PyTorch model.
"""
from typing import Dict, List, Tuple

import torch


def train_step(model: torch.nn.Module,
               dataloader: torch.utils.data.DataLoader,
               loss_fn: torch.nn.Module,
               optimizer: torch.optim.Optimizer,
               device: torch.device) -> Tuple[float, float]:
    """Trains a PyTorch model for a single epoch.

    Puts the model in training mode, then runs the training steps
    (forward pass, loss, zero grad, backward pass, optimizer step)
    on every batch of the DataLoader.

    Args:
        model: The PyTorch model to train.
        dataloader: DataLoader with the training data.
        loss_fn: Loss function to minimize.
        optimizer: Optimizer that updates the model's parameters.
        device: Device to compute on, e.g. "cuda" or "cpu".

    Returns:
        A tuple (train_loss, train_accuracy), each averaged over the batches,
        with accuracy as a fraction between 0 and 1, e.g. (0.1112, 0.8743).
    """
    model.train()
    train_loss, train_acc = 0, 0

    for X, y in dataloader:
        X, y = X.to(device), y.to(device)

        # 1. Forward pass
        y_pred = model(X)

        # 2. Calculate and accumulate the loss
        loss = loss_fn(y_pred, y)
        train_loss += loss.item()

        # 3. Zero the gradients
        optimizer.zero_grad()

        # 4. Backpropagate the loss
        loss.backward()

        # 5. Update the parameters
        optimizer.step()

        # Accuracy for this batch: fraction of predicted classes that are right
        y_pred_class = torch.softmax(y_pred, dim=1).argmax(dim=1)
        train_acc += (y_pred_class == y).sum().item() / len(y_pred)

    # Average over the number of batches
    train_loss = train_loss / len(dataloader)
    train_acc = train_acc / len(dataloader)
    return train_loss, train_acc


def test_step(model: torch.nn.Module,
              dataloader: torch.utils.data.DataLoader,
              loss_fn: torch.nn.Module,
              device: torch.device) -> Tuple[float, float]:
    """Tests a PyTorch model for a single epoch.

    Puts the model in evaluation mode and runs a forward pass over
    every batch of the DataLoader, without updating any parameters.

    Args:
        model: The PyTorch model to test.
        dataloader: DataLoader with the test data.
        loss_fn: Loss function to evaluate the predictions with.
        device: Device to compute on, e.g. "cuda" or "cpu".

    Returns:
        A tuple (test_loss, test_accuracy), each averaged over the batches,
        with accuracy as a fraction between 0 and 1, e.g. (0.0223, 0.8985).
    """
    model.eval()
    test_loss, test_acc = 0, 0

    with torch.inference_mode():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)

            # 1. Forward pass
            test_pred_logits = model(X)

            # 2. Calculate and accumulate the loss
            loss = loss_fn(test_pred_logits, y)
            test_loss += loss.item()

            # Accuracy for this batch
            test_pred_labels = torch.softmax(test_pred_logits, dim=1).argmax(dim=1)
            test_acc += (test_pred_labels == y).sum().item() / len(test_pred_labels)

    test_loss = test_loss / len(dataloader)
    test_acc = test_acc / len(dataloader)
    return test_loss, test_acc


def train(model: torch.nn.Module,
          train_dataloader: torch.utils.data.DataLoader,
          test_dataloader: torch.utils.data.DataLoader,
          optimizer: torch.optim.Optimizer,
          loss_fn: torch.nn.Module,
          epochs: int,
          device: torch.device) -> Dict[str, List[float]]:
    """Trains and tests a PyTorch model.

    Runs train_step() and test_step() once per epoch, prints the
    metrics, and records them.

    Args:
        model: The PyTorch model to train and test.
        train_dataloader: DataLoader with the training data.
        test_dataloader: DataLoader with the test data.
        optimizer: Optimizer that updates the model's parameters.
        loss_fn: Loss function used on both datasets.
        epochs: Number of epochs to train for.
        device: Device to compute on, e.g. "cuda" or "cpu".

    Returns:
        A dictionary of lists with one value per epoch:
        {"train_loss": [...], "train_acc": [...],
         "test_loss": [...], "test_acc": [...]}
    """
    results = {"train_loss": [], "train_acc": [], "test_loss": [], "test_acc": []}

    model.to(device)

    for epoch in range(epochs):
        train_loss, train_acc = train_step(model=model,
                                           dataloader=train_dataloader,
                                           loss_fn=loss_fn,
                                           optimizer=optimizer,
                                           device=device)
        test_loss, test_acc = test_step(model=model,
                                        dataloader=test_dataloader,
                                        loss_fn=loss_fn,
                                        device=device)

        print(
            f"Epoch: {epoch+1} | "
            f"train_loss: {train_loss:.4f} | "
            f"train_acc: {train_acc:.4f} | "
            f"test_loss: {test_loss:.4f} | "
            f"test_acc: {test_acc:.4f}"
        )

        results["train_loss"].append(train_loss)
        results["train_acc"].append(train_acc)
        results["test_loss"].append(test_loss)
        results["test_acc"].append(test_acc)

    return results

%%writefile going_modular/utils.py
"""
Utility functions for PyTorch model training and saving.
"""
from pathlib import Path

import torch


def save_model(model: torch.nn.Module,
               target_dir: str,
               model_name: str):
    """Saves a PyTorch model's state_dict to a target folder.

    Args:
        model: The PyTorch model to save.
        target_dir: Folder to save the model in (created if missing).
        model_name: File name for the saved model; must end in ".pth" or ".pt".

    Example:
        save_model(model=model_0,
                   target_dir="models",
                   model_name="05_going_modular_tinyvgg_model.pth")
    """
    # Create the target folder
    target_dir_path = Path(target_dir)
    target_dir_path.mkdir(parents=True, exist_ok=True)

    # Build the save path
    assert model_name.endswith(".pth") or model_name.endswith(".pt"), \
        "model_name should end with '.pt' or '.pth'"
    model_save_path = target_dir_path / model_name

    # Save only the learned parameters
    print(f"[INFO] Saving model to: {model_save_path}")
    torch.save(obj=model.state_dict(), f=model_save_path)

%%writefile going_modular/train.py
"""
Trains a TinyVGG image classification model using device-agnostic code.
"""
import torch
from torchvision import transforms

import data_setup, engine, model_builder, utils

# Hyperparameters
NUM_EPOCHS = 5
BATCH_SIZE = 32
HIDDEN_UNITS = 10
LEARNING_RATE = 0.001

# Data folders
train_dir = "data/pizza_steak_sushi/train"
test_dir = "data/pizza_steak_sushi/test"

# Target device
device = "cuda" if torch.cuda.is_available() else "cpu"

# Transforms: resize to 64x64 (what TinyVGG expects) and convert to a tensor
data_transform = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor(),
])

# DataLoaders, from data_setup.py
train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(
    train_dir=train_dir,
    test_dir=test_dir,
    transform=data_transform,
    batch_size=BATCH_SIZE,
)

# Model, from model_builder.py
torch.manual_seed(42)
model = model_builder.TinyVGG(
    input_shape=3,
    hidden_units=HIDDEN_UNITS,
    output_shape=len(class_names),
).to(device)

# Loss function and optimizer
loss_fn = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)

# Training, with engine.py
engine.train(model=model,
             train_dataloader=train_dataloader,
             test_dataloader=test_dataloader,
             loss_fn=loss_fn,
             optimizer=optimizer,
             epochs=NUM_EPOCHS,
             device=device)

# Save the model, with utils.py
utils.save_model(model=model,
                 target_dir="models",
                 model_name="05_going_modular_script_mode_tinyvgg_model.pth")

!python going_modular/train.py

import torch
from torchvision import transforms

from going_modular import data_setup, engine, model_builder

device = "cuda" if torch.cuda.is_available() else "cpu"

data_transform = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor(),
])

train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(
    train_dir="data/pizza_steak_sushi/train",
    test_dir="data/pizza_steak_sushi/test",
    transform=data_transform,
    batch_size=32,
)
class_names

loaded_model = model_builder.TinyVGG(input_shape=3,
                                     hidden_units=10,
                                     output_shape=len(class_names))
loaded_model.load_state_dict(
    torch.load("models/05_going_modular_script_mode_tinyvgg_model.pth", map_location=device)
)
loaded_model.to(device)

test_loss, test_acc = engine.test_step(model=loaded_model,
                                       dataloader=test_dataloader,
                                       loss_fn=torch.nn.CrossEntropyLoss(),
                                       device=device)
print(f"Loaded model | test_loss: {test_loss:.4f} | test_acc: {test_acc:.4f}")

%%writefile going_modular/train.py
"""
Trains a TinyVGG image classification model using device-agnostic code.

Example:
    python going_modular/train.py --num_epochs 10 --learning_rate 0.003
"""
import argparse

import torch
from torchvision import transforms

import data_setup, engine, model_builder, utils

# 1. Describe the arguments the script accepts
parser = argparse.ArgumentParser(description="Train TinyVGG on pizza, steak and sushi images.")
parser.add_argument("--num_epochs", type=int, default=5,
                    help="number of epochs to train for (default: 5)")
parser.add_argument("--batch_size", type=int, default=32,
                    help="number of images per batch (default: 32)")
parser.add_argument("--hidden_units", type=int, default=10,
                    help="number of filters in each TinyVGG layer (default: 10)")
parser.add_argument("--learning_rate", type=float, default=0.001,
                    help="learning rate for the Adam optimizer (default: 0.001)")

# 2. Read what was typed on the command line
args = parser.parse_args()

# 3. Use the values
NUM_EPOCHS = args.num_epochs
BATCH_SIZE = args.batch_size
HIDDEN_UNITS = args.hidden_units
LEARNING_RATE = args.learning_rate
print(f"[INFO] Training for {NUM_EPOCHS} epochs | batch size {BATCH_SIZE} | "
      f"{HIDDEN_UNITS} hidden units | learning rate {LEARNING_RATE}")

# Everything below is unchanged
train_dir = "data/pizza_steak_sushi/train"
test_dir = "data/pizza_steak_sushi/test"

device = "cuda" if torch.cuda.is_available() else "cpu"

data_transform = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor(),
])

train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(
    train_dir=train_dir,
    test_dir=test_dir,
    transform=data_transform,
    batch_size=BATCH_SIZE,
)

torch.manual_seed(42)
model = model_builder.TinyVGG(
    input_shape=3,
    hidden_units=HIDDEN_UNITS,
    output_shape=len(class_names),
).to(device)

loss_fn = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)

engine.train(model=model,
             train_dataloader=train_dataloader,
             test_dataloader=test_dataloader,
             loss_fn=loss_fn,
             optimizer=optimizer,
             epochs=NUM_EPOCHS,
             device=device)

utils.save_model(model=model,
                 target_dir="models",
                 model_name="05_going_modular_script_mode_tinyvgg_model.pth")

!python going_modular/train.py --help

!python going_modular/train.py --num_epochs 3 --hidden_units 20 --learning_rate 0.003