init

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 import os
 import torch
 
-from model import create_vit16_model
+from model import create_model
 from timeit import default_timer as timer
 from typing import Tuple, Dict
 
@@ -14,14 +14,14 @@ with open("class_names.txt", "r") as f: # reading them in from class_names.txt
 ### 2. Model and transforms preparation ###    
 
 # Create model
-vit16, vit16_transforms = create_vit16_model(
-    num_classes=101, # could also use len(class_names)
+model_created, model_transforms = create_model(
+    num_classes=len(class_names),
 )
 
 
-vit16.load_state_dict(
+model_created.load_state_dict(
     torch.load(
-        f="model_food101_20_percent.pth",
+        f="model_1.pth",
         map_location=torch.device("cpu"),  # load to CPU
     )
 )
@@ -38,13 +38,13 @@ def predict(img) -> Tuple[Dict, float]:
     start_time = timer()
     
     # Transform the target image and add a batch dimension
-    img = vit16_transforms(img).unsqueeze(0)
+    img = model_transforms(img).unsqueeze(0)
     
     # Put model into evaluation mode and turn on inference mode
-    vit16.eval()
+    model_created.eval()
     with torch.inference_mode():
         # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
-        pred_probs = torch.softmax(vit16(img), dim=1)
+        pred_probs = torch.softmax(model_created(img), dim=1)
     
     # Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
     pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
@@ -58,8 +58,8 @@ def predict(img) -> Tuple[Dict, float]:
 ### 4. Gradio app ###
 
 # Create title, description and article strings
-title = "FoodVision ViT 🍔👁"
-description = "A ViT_B_16 feature extractor computer vision model to classify images of food into 101 different classes using 20% of the data."
+title = "World Puzzle Solver"
+description = "A World Puzzle Solver app that uses a PyTorch model to predict the letters in a target image."
 article = ""
 
 # Create examples list from "examples/" directory

class_names.txt

@@ -1,101 +1,32 @@
-apple_pie
-baby_back_ribs
-baklava
-beef_carpaccio
-beef_tartare
-beet_salad
-beignets
-bibimbap
-bread_pudding
-breakfast_burrito
-bruschetta
-caesar_salad
-cannoli
-caprese_salad
-carrot_cake
-ceviche
-cheese_plate
-cheesecake
-chicken_curry
-chicken_quesadilla
-chicken_wings
-chocolate_cake
-chocolate_mousse
-churros
-clam_chowder
-club_sandwich
-crab_cakes
-creme_brulee
-croque_madame
-cup_cakes
-deviled_eggs
-donuts
-dumplings
-edamame
-eggs_benedict
-escargots
-falafel
-filet_mignon
-fish_and_chips
-foie_gras
-french_fries
-french_onion_soup
-french_toast
-fried_calamari
-fried_rice
-frozen_yogurt
-garlic_bread
-gnocchi
-greek_salad
-grilled_cheese_sandwich
-grilled_salmon
-guacamole
-gyoza
-hamburger
-hot_and_sour_soup
-hot_dog
-huevos_rancheros
-hummus
-ice_cream
-lasagna
-lobster_bisque
-lobster_roll_sandwich
-macaroni_and_cheese
-macarons
-miso_soup
-mussels
-nachos
-omelette
-onion_rings
-oysters
-pad_thai
-paella
-pancakes
-panna_cotta
-peking_duck
-pho
-pizza
-pork_chop
-poutine
-prime_rib
-pulled_pork_sandwich
-ramen
-ravioli
-red_velvet_cake
-risotto
-samosa
-sashimi
-scallops
-seaweed_salad
-shrimp_and_grits
-spaghetti_bolognese
-spaghetti_carbonara
-spring_rolls
-steak
-strawberry_shortcake
-sushi
-tacos
-takoyaki
-tiramisu
-tuna_tartare
-waffles
+A
+B
+C
+D
+E
+F
+G
+H
+I
+J
+K
+L
+M
+N
+O
+P
+Q
+R
+S
+T
+U
+V
+W
+X
+Y
+Z
+Á
+É
+Í
+Ñ
+Ó
+Ú

helper/data_setup.py

@@ -0,0 +1,66 @@
+"""
+Contains functionality for creating PyTorch DataLoaders for 
+image classification data.
+"""
+import os
+
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+
+NUM_WORKERS = os.cpu_count()
+
+def create_dataloaders(
+    train_dir: str, 
+    test_dir: str, 
+    train_transform: transforms.Compose,
+    test_transform: transforms.Compose, 
+    batch_size: int, 
+    num_workers: int=NUM_WORKERS
+):
+  """Creates training and testing DataLoaders.
+
+  Takes in a training directory and testing directory path and turns
+  them into PyTorch Datasets and then into PyTorch DataLoaders.
+
+  Args:
+    train_dir: Path to training directory.
+    test_dir: Path to testing directory.
+    transform: torchvision transforms to perform on training and testing data.
+    batch_size: Number of samples per batch in each of the DataLoaders.
+    num_workers: An integer for number of workers per DataLoader.
+
+  Returns:
+    A tuple of (train_dataloader, test_dataloader, class_names).
+    Where class_names is a list of the target classes.
+    Example usage:
+      train_dataloader, test_dataloader, class_names = \
+        = create_dataloaders(train_dir=path/to/train_dir,
+                             test_dir=path/to/test_dir,
+                             transform=some_transform,
+                             batch_size=32,
+                             num_workers=4)
+  """
+  # Use ImageFolder to create dataset(s)
+  train_data = datasets.ImageFolder(train_dir, transform=train_transform)
+  test_data = datasets.ImageFolder(test_dir, transform=test_transform)
+
+  # Get class names
+  class_names = train_data.classes
+
+  # Turn images into data loaders
+  train_dataloader = DataLoader(
+      train_data,
+      batch_size=batch_size,
+      shuffle=True,
+      num_workers=num_workers,
+      pin_memory=True,
+  )
+  test_dataloader = DataLoader(
+      test_data,
+      batch_size=batch_size,
+      shuffle=False,
+      num_workers=num_workers,
+      pin_memory=True,
+  )
+
+  return train_dataloader, test_dataloader, class_names

helper/engine.py

@@ -0,0 +1,195 @@
+"""
+Contains functions for training and testing a PyTorch model.
+"""
+import torch
+
+from tqdm.auto import tqdm
+from typing import Dict, List, Tuple
+
+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.
+
+    Turns a target PyTorch model to training mode and then
+    runs through all of the required training steps (forward
+    pass, loss calculation, optimizer step).
+
+    Args:
+    model: A PyTorch model to be trained.
+    dataloader: A DataLoader instance for the model to be trained on.
+    loss_fn: A PyTorch loss function to minimize.
+    optimizer: A PyTorch optimizer to help minimize the loss function.
+    device: A target device to compute on (e.g. "cuda" or "cpu").
+
+    Returns:
+    A tuple of training loss and training accuracy metrics.
+    In the form (train_loss, train_accuracy). For example:
+
+    (0.1112, 0.8743)
+    """
+    # Put model in train mode
+    model.train()
+
+    # Setup train loss and train accuracy values
+    train_loss, train_acc = 0, 0
+
+    # Loop through data loader data batches
+    for batch, (X, y) in enumerate(dataloader):
+        # Send data to target device
+        X, y = X.to(device), y.to(device)
+
+        # 1. Forward pass
+        y_pred = model(X)
+
+        # 2. Calculate  and accumulate loss
+        loss = loss_fn(y_pred, y)
+        train_loss += loss.item() 
+
+        # 3. Optimizer zero grad
+        optimizer.zero_grad()
+
+        # 4. Loss backward
+        loss.backward()
+
+        # 5. Optimizer step
+        optimizer.step()
+
+        # Calculate and accumulate accuracy metric across all batches
+        y_pred_class = torch.argmax(torch.softmax(y_pred, dim=1), dim=1)
+        train_acc += (y_pred_class == y).sum().item()/len(y_pred)
+
+    # Adjust metrics to get average loss and accuracy per batch 
+    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.
+
+    Turns a target PyTorch model to "eval" mode and then performs
+    a forward pass on a testing dataset.
+
+    Args:
+    model: A PyTorch model to be tested.
+    dataloader: A DataLoader instance for the model to be tested on.
+    loss_fn: A PyTorch loss function to calculate loss on the test data.
+    device: A target device to compute on (e.g. "cuda" or "cpu").
+
+    Returns:
+    A tuple of testing loss and testing accuracy metrics.
+    In the form (test_loss, test_accuracy). For example:
+
+    (0.0223, 0.8985)
+    """
+    # Put model in eval mode
+    model.eval() 
+
+    # Setup test loss and test accuracy values
+    test_loss, test_acc = 0, 0
+
+    # Turn on inference context manager
+    with torch.inference_mode():
+        # Loop through DataLoader batches
+        for batch, (X, y) in enumerate(dataloader):
+            # Send data to target device
+            X, y = X.to(device), y.to(device)
+
+            # 1. Forward pass
+            test_pred_logits = model(X)
+
+            # 2. Calculate and accumulate loss
+            loss = loss_fn(test_pred_logits, y)
+            test_loss += loss.item()
+
+            # Calculate and accumulate accuracy
+            test_pred_labels = test_pred_logits.argmax(dim=1)
+            test_acc += ((test_pred_labels == y).sum().item()/len(test_pred_labels))
+
+    # Adjust metrics to get average loss and accuracy per batch 
+    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]:
+    """Trains and tests a PyTorch model.
+
+    Passes a target PyTorch models through train_step() and test_step()
+    functions for a number of epochs, training and testing the model
+    in the same epoch loop.
+
+    Calculates, prints and stores evaluation metrics throughout.
+
+    Args:
+    model: A PyTorch model to be trained and tested.
+    train_dataloader: A DataLoader instance for the model to be trained on.
+    test_dataloader: A DataLoader instance for the model to be tested on.
+    optimizer: A PyTorch optimizer to help minimize the loss function.
+    loss_fn: A PyTorch loss function to calculate loss on both datasets.
+    epochs: An integer indicating how many epochs to train for.
+    device: A target device to compute on (e.g. "cuda" or "cpu").
+
+    Returns:
+    A dictionary of training and testing loss as well as training and
+    testing accuracy metrics. Each metric has a value in a list for 
+    each epoch.
+    In the form: {train_loss: [...],
+              train_acc: [...],
+              test_loss: [...],
+              test_acc: [...]} 
+    For example if training for epochs=2: 
+             {train_loss: [2.0616, 1.0537],
+              train_acc: [0.3945, 0.3945],
+              test_loss: [1.2641, 1.5706],
+              test_acc: [0.3400, 0.2973]} 
+    """
+    # Create empty results dictionary
+    results = {"train_loss": [],
+               "train_acc": [],
+               "test_loss": [],
+               "test_acc": []
+    }
+    
+    # Make sure model on target device
+    model.to(device)
+
+    # Loop through training and testing steps for a number of epochs
+    for epoch in tqdm(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 out what's happening
+        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}"
+        )
+
+        # Update results dictionary
+        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 the filled results at the end of the epochs
+    return results

helper/helper_functions.py

@@ -0,0 +1,294 @@
+"""
+A series of helper functions used throughout the course.
+
+If a function gets defined once and could be used over and over, it'll go in here.
+"""
+import torch
+import matplotlib.pyplot as plt
+import numpy as np
+
+from torch import nn
+
+import os
+import zipfile
+
+from pathlib import Path
+
+import requests
+
+# Walk through an image classification directory and find out how many files (images)
+# are in each subdirectory.
+import os
+
+def walk_through_dir(dir_path):
+    """
+    Walks through dir_path returning its contents.
+    Args:
+    dir_path (str): target directory
+
+    Returns:
+    A print out of:
+      number of subdiretories in dir_path
+      number of images (files) in each subdirectory
+      name of each subdirectory
+    """
+    for dirpath, dirnames, filenames in os.walk(dir_path):
+        print(f"There are {len(dirnames)} directories and {len(filenames)} images in '{dirpath}'.")
+
+def plot_decision_boundary(model: torch.nn.Module, X: torch.Tensor, y: torch.Tensor):
+    """Plots decision boundaries of model predicting on X in comparison to y.
+
+    Source - https://madewithml.com/courses/foundations/neural-networks/ (with modifications)
+    """
+    # Put everything to CPU (works better with NumPy + Matplotlib)
+    model.to("cpu")
+    X, y = X.to("cpu"), y.to("cpu")
+
+    # Setup prediction boundaries and grid
+    x_min, x_max = X[:, 0].min() - 0.1, X[:, 0].max() + 0.1
+    y_min, y_max = X[:, 1].min() - 0.1, X[:, 1].max() + 0.1
+    xx, yy = np.meshgrid(np.linspace(x_min, x_max, 101), np.linspace(y_min, y_max, 101))
+
+    # Make features
+    X_to_pred_on = torch.from_numpy(np.column_stack((xx.ravel(), yy.ravel()))).float()
+
+    # Make predictions
+    model.eval()
+    with torch.inference_mode():
+        y_logits = model(X_to_pred_on)
+
+    # Test for multi-class or binary and adjust logits to prediction labels
+    if len(torch.unique(y)) > 2:
+        y_pred = torch.softmax(y_logits, dim=1).argmax(dim=1)  # mutli-class
+    else:
+        y_pred = torch.round(torch.sigmoid(y_logits))  # binary
+
+    # Reshape preds and plot
+    y_pred = y_pred.reshape(xx.shape).detach().numpy()
+    plt.contourf(xx, yy, y_pred, cmap=plt.cm.RdYlBu, alpha=0.7)
+    plt.scatter(X[:, 0], X[:, 1], c=y, s=40, cmap=plt.cm.RdYlBu)
+    plt.xlim(xx.min(), xx.max())
+    plt.ylim(yy.min(), yy.max())
+
+
+# Plot linear data or training and test and predictions (optional)
+def plot_predictions(
+    train_data, train_labels, test_data, test_labels, predictions=None
+):
+    """
+  Plots linear training data and test data and compares predictions.
+  """
+    plt.figure(figsize=(10, 7))
+
+    # Plot training data in blue
+    plt.scatter(train_data, train_labels, c="b", s=4, label="Training data")
+
+    # Plot test data in green
+    plt.scatter(test_data, test_labels, c="g", s=4, label="Testing data")
+
+    if predictions is not None:
+        # Plot the predictions in red (predictions were made on the test data)
+        plt.scatter(test_data, predictions, c="r", s=4, label="Predictions")
+
+    # Show the legend
+    plt.legend(prop={"size": 14})
+
+
+# Calculate accuracy (a classification metric)
+def accuracy_fn(y_true, y_pred):
+    """Calculates accuracy between truth labels and predictions.
+
+    Args:
+        y_true (torch.Tensor): Truth labels for predictions.
+        y_pred (torch.Tensor): Predictions to be compared to predictions.
+
+    Returns:
+        [torch.float]: Accuracy value between y_true and y_pred, e.g. 78.45
+    """
+    correct = torch.eq(y_true, y_pred).sum().item()
+    acc = (correct / len(y_pred)) * 100
+    return acc
+
+
+def print_train_time(start, end, device=None):
+    """Prints difference between start and end time.
+
+    Args:
+        start (float): Start time of computation (preferred in timeit format). 
+        end (float): End time of computation.
+        device ([type], optional): Device that compute is running on. Defaults to None.
+
+    Returns:
+        float: time between start and end in seconds (higher is longer).
+    """
+    total_time = end - start
+    print(f"\nTrain time on {device}: {total_time:.3f} seconds")
+    return total_time
+
+
+# Plot loss curves of a model
+def plot_loss_curves(results):
+    """Plots training curves of a results dictionary.
+
+    Args:
+        results (dict): dictionary containing list of values, e.g.
+            {"train_loss": [...],
+             "train_acc": [...],
+             "test_loss": [...],
+             "test_acc": [...]}
+    """
+    loss = results["train_loss"]
+    test_loss = results["test_loss"]
+
+    accuracy = results["train_acc"]
+    test_accuracy = results["test_acc"]
+
+    epochs = range(len(results["train_loss"]))
+
+    plt.figure(figsize=(15, 7))
+
+    # Plot loss
+    plt.subplot(1, 2, 1)
+    plt.plot(epochs, loss, label="train_loss")
+    plt.plot(epochs, test_loss, label="test_loss")
+    plt.title("Loss")
+    plt.xlabel("Epochs")
+    plt.legend()
+
+    # Plot accuracy
+    plt.subplot(1, 2, 2)
+    plt.plot(epochs, accuracy, label="train_accuracy")
+    plt.plot(epochs, test_accuracy, label="test_accuracy")
+    plt.title("Accuracy")
+    plt.xlabel("Epochs")
+    plt.legend()
+
+
+# Pred and plot image function from notebook 04
+# See creation: https://www.learnpytorch.io/04_pytorch_custom_datasets/#113-putting-custom-image-prediction-together-building-a-function
+from typing import List
+import torchvision
+
+
+def pred_and_plot_image(
+    model: torch.nn.Module,
+    image_path: str,
+    class_names: List[str] = None,
+    transform=None,
+    device: torch.device = "cuda" if torch.cuda.is_available() else "cpu",
+):
+    """Makes a prediction on a target image with a trained model and plots the image.
+
+    Args:
+        model (torch.nn.Module): trained PyTorch image classification model.
+        image_path (str): filepath to target image.
+        class_names (List[str], optional): different class names for target image. Defaults to None.
+        transform (_type_, optional): transform of target image. Defaults to None.
+        device (torch.device, optional): target device to compute on. Defaults to "cuda" if torch.cuda.is_available() else "cpu".
+    
+    Returns:
+        Matplotlib plot of target image and model prediction as title.
+
+    Example usage:
+        pred_and_plot_image(model=model,
+                            image="some_image.jpeg",
+                            class_names=["class_1", "class_2", "class_3"],
+                            transform=torchvision.transforms.ToTensor(),
+                            device=device)
+    """
+
+    # 1. Load in image and convert the tensor values to float32
+    target_image = torchvision.io.read_image(str(image_path)).type(torch.float32)
+
+    # 2. Divide the image pixel values by 255 to get them between [0, 1]
+    target_image = target_image / 255.0
+
+    # 3. Transform if necessary
+    if transform:
+        target_image = transform(target_image)
+
+    # 4. Make sure the model is on the target device
+    model.to(device)
+
+    # 5. Turn on model evaluation mode and inference mode
+    model.eval()
+    with torch.inference_mode():
+        # Add an extra dimension to the image
+        target_image = target_image.unsqueeze(dim=0)
+
+        # Make a prediction on image with an extra dimension and send it to the target device
+        target_image_pred = model(target_image.to(device))
+
+    # 6. Convert logits -> prediction probabilities (using torch.softmax() for multi-class classification)
+    target_image_pred_probs = torch.softmax(target_image_pred, dim=1)
+
+    # 7. Convert prediction probabilities -> prediction labels
+    target_image_pred_label = torch.argmax(target_image_pred_probs, dim=1)
+
+    # 8. Plot the image alongside the prediction and prediction probability
+    plt.imshow(
+        target_image.squeeze().permute(1, 2, 0)
+    )  # make sure it's the right size for matplotlib
+    if class_names:
+        title = f"Pred: {class_names[target_image_pred_label.cpu()]} | Prob: {target_image_pred_probs.max().cpu():.3f}"
+    else:
+        title = f"Pred: {target_image_pred_label} | Prob: {target_image_pred_probs.max().cpu():.3f}"
+    plt.title(title)
+    plt.axis(False)
+
+def set_seeds(seed: int=42):
+    """Sets random sets for torch operations.
+
+    Args:
+        seed (int, optional): Random seed to set. Defaults to 42.
+    """
+    # Set the seed for general torch operations
+    torch.manual_seed(seed)
+    # Set the seed for CUDA torch operations (ones that happen on the GPU)
+    torch.cuda.manual_seed(seed)
+
+def download_data(source: str, 
+                  destination: str,
+                  remove_source: bool = True) -> Path:
+    """Downloads a zipped dataset from source and unzips to destination.
+
+    Args:
+        source (str): A link to a zipped file containing data.
+        destination (str): A target directory to unzip data to.
+        remove_source (bool): Whether to remove the source after downloading and extracting.
+    
+    Returns:
+        pathlib.Path to downloaded data.
+    
+    Example usage:
+        download_data(source="https://github.com/mrdbourke/pytorch-deep-learning/raw/main/data/pizza_steak_sushi.zip",
+                      destination="pizza_steak_sushi")
+    """
+    # Setup path to data folder
+    data_path = Path("data/")
+    image_path = data_path / destination
+
+    # If the image folder doesn't exist, download it and prepare it... 
+    if image_path.is_dir():
+        print(f"[INFO] {image_path} directory exists, skipping download.")
+    else:
+        print(f"[INFO] Did not find {image_path} directory, creating one...")
+        image_path.mkdir(parents=True, exist_ok=True)
+        
+        # Download pizza, steak, sushi data
+        target_file = Path(source).name
+        with open(data_path / target_file, "wb") as f:
+            request = requests.get(source)
+            print(f"[INFO] Downloading {target_file} from {source}...")
+            f.write(request.content)
+
+        # Unzip pizza, steak, sushi data
+        with zipfile.ZipFile(data_path / target_file, "r") as zip_ref:
+            print(f"[INFO] Unzipping {target_file} data...") 
+            zip_ref.extractall(image_path)
+
+        # Remove .zip file
+        if remove_source:
+            os.remove(data_path / target_file)
+    
+    return image_path

helper/model_builder.py

@@ -0,0 +1,56 @@
+"""
+Contains PyTorch model code to instantiate a TinyVGG model.
+"""
+import torch
+from torch import nn 
+
+class TinyVGG(nn.Module):
+    """Creates the TinyVGG architecture.
+
+    Replicates the TinyVGG architecture from the CNN explainer website in PyTorch.
+    See the original architecture here: https://poloclub.github.io/cnn-explainer/
+
+    Args:
+    input_shape: An integer indicating number of input channels.
+    hidden_units: An integer indicating number of hidden units between layers.
+    output_shape: An integer indicating number of output units.
+    """
+    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(),
+          # Where did this in_features shape come from? 
+          # It's because each layer of our network compresses and changes the shape of our inputs data.
+          nn.Linear(in_features=hidden_units*13*13,
+                    out_features=output_shape)
+        )
+    
+    def forward(self, x: torch.Tensor):
+        x = self.conv_block_1(x)
+        x = self.conv_block_2(x)
+        x = self.classifier(x)
+        return x
+        # return self.classifier(self.block_2(self.block_1(x))) # <- leverage the benefits of operator fusion

helper/predictions.py

@@ -0,0 +1,83 @@
+"""
+Utility functions to make predictions.
+
+Main reference for code creation: https://www.learnpytorch.io/06_pytorch_transfer_learning/#6-make-predictions-on-images-from-the-test-set 
+"""
+import torch
+import torchvision
+from torchvision import transforms
+import matplotlib.pyplot as plt
+
+from typing import List, Tuple
+
+from PIL import Image
+
+# Set device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Predict on a target image with a target model
+# Function created in: https://www.learnpytorch.io/06_pytorch_transfer_learning/#6-make-predictions-on-images-from-the-test-set
+def pred_and_plot_image(
+    model: torch.nn.Module,
+    class_names: List[str],
+    image_path: str,
+    image_size: Tuple[int, int] = (224, 224),
+    transform: torchvision.transforms = None,
+    device: torch.device = device,
+):
+    """Predicts on a target image with a target model.
+
+    Args:
+        model (torch.nn.Module): A trained (or untrained) PyTorch model to predict on an image.
+        class_names (List[str]): A list of target classes to map predictions to.
+        image_path (str): Filepath to target image to predict on.
+        image_size (Tuple[int, int], optional): Size to transform target image to. Defaults to (224, 224).
+        transform (torchvision.transforms, optional): Transform to perform on image. Defaults to None which uses ImageNet normalization.
+        device (torch.device, optional): Target device to perform prediction on. Defaults to device.
+    """
+
+    # Open image
+    img = Image.open(image_path)
+
+    # Create transformation for image (if one doesn't exist)
+    if transform is not None:
+        image_transform = transform
+    else:
+        image_transform = transforms.Compose(
+            [
+                transforms.Resize(image_size),
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+                ),
+            ]
+        )
+
+    ### Predict on image ###
+
+    # Make sure the model is on the target device
+    model.to(device)
+
+    # Turn on model evaluation mode and inference mode
+    model.eval()
+    with torch.inference_mode():
+        # Transform and add an extra dimension to image (model requires samples in [batch_size, color_channels, height, width])
+        transformed_image = image_transform(img).unsqueeze(dim=0)
+
+        # Make a prediction on image with an extra dimension and send it to the target device
+        target_image_pred = model(transformed_image.to(device))
+
+    # Convert logits -> prediction probabilities (using torch.softmax() for multi-class classification)
+    target_image_pred_probs = torch.softmax(target_image_pred, dim=1)
+
+    # Convert prediction probabilities -> prediction labels
+    target_image_pred_label = torch.argmax(target_image_pred_probs, dim=1)
+
+    # Plot image with predicted label and probability
+    plt.figure()
+    plt.imshow(img)
+    plt.title(
+        f"Pred: {class_names[target_image_pred_label]} | Prob: {target_image_pred_probs.max():.3f}"
+    )
+    plt.axis(False)
+

helper/train.py

@@ -0,0 +1,62 @@
+"""
+Trains a PyTorch image classification model using device-agnostic code.
+"""
+
+import os
+import torch
+import data_setup, engine, model_builder, utils
+
+from torchvision import transforms
+
+# Setup hyperparameters
+NUM_EPOCHS = 5
+BATCH_SIZE = 32
+HIDDEN_UNITS = 10
+LEARNING_RATE = 0.001
+
+# Setup directories
+train_dir = "data/pizza_steak_sushi/train"
+test_dir = "data/pizza_steak_sushi/test"
+
+# Setup target device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Create transforms
+data_transform = transforms.Compose([
+  transforms.Resize((64, 64)),
+  transforms.ToTensor()
+])
+
+# Create DataLoaders with help 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
+)
+
+# Create model with help from model_builder.py
+model = model_builder.TinyVGG(
+    input_shape=3,
+    hidden_units=HIDDEN_UNITS,
+    output_shape=len(class_names)
+).to(device)
+
+# Set loss and optimizer
+loss_fn = torch.nn.CrossEntropyLoss()
+optimizer = torch.optim.Adam(model.parameters(),
+                             lr=LEARNING_RATE)
+
+# Start training with help from 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 help from utils.py
+utils.save_model(model=model,
+                 target_dir="models",
+                 model_name="05_going_modular_script_mode_tinyvgg_model.pth")

helper/utils.py

@@ -0,0 +1,35 @@
+"""
+Contains various utility functions for PyTorch model training and saving.
+"""
+import torch
+from pathlib import Path
+
+def save_model(model: torch.nn.Module,
+               target_dir: str,
+               model_name: str):
+    """Saves a PyTorch model to a target directory.
+
+    Args:
+    model: A target PyTorch model to save.
+    target_dir: A directory for saving the model to.
+    model_name: A filename for the saved model. Should include
+      either ".pth" or ".pt" as the file extension.
+
+    Example usage:
+    save_model(model=model_0,
+               target_dir="models",
+               model_name="05_going_modular_tingvgg_model.pth")
+    """
+    # Create target directory
+    target_dir_path = Path(target_dir)
+    target_dir_path.mkdir(parents=True,
+                        exist_ok=True)
+
+    # Create model 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 the model state_dict()
+    print(f"[INFO] Saving model to: {model_save_path}")
+    torch.save(obj=model.state_dict(),
+             f=model_save_path)

model.py

@@ -1,34 +1,67 @@
 import torch
 import torchvision
-
 from torch import nn
 
 
-def create_vit16_model(num_classes:int=101, 
-                          seed:int=42):
-    """Creates an vit16 feature extractor model and transforms.
+def create_model(num_classes: int = 32,
+                          seed: int = 42):
+    """Creates a feature extractor model and transforms.
 
     Args:
-        num_classes (int, optional): number of classes in the classifier head. 
-            Defaults to 3.
+        num_classes (int, optional): number of classes in the classifier head.
+            Defaults to 32.
         seed (int, optional): random seed value. Defaults to 42.
 
     Returns:
-        model (torch.nn.Module): vit feature extractor model. 
+        model (torch.nn.Module): vit feature extractor model.
         transforms (torchvision.transforms): vit image transforms.
     """
-    # Create vit pretrained weights, transforms and model
-    weights = torchvision.models.ViT_B_16_Weights.DEFAULT;
-    transforms = weights.transforms()
-    model = torchvision.models.vit_b_16(weights=weights)
+    IMG_SIZE = 28
+    transforms = transforms.Compose([
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.Grayscale(num_output_channels=1),
+        transforms.ToTensor()])
 
-    # Freeze all layers in base model
-    for param in model.parameters():
-        param.requires_grad = False
+        # Create a convolutional neural network
+    class Model(nn.Module):
+            def __init__(self, input_shape: int, hidden_units: int, output_shape: int):
+                super().__init__()
+                self.block_1 = nn.Sequential(
+                        nn.Conv2d(in_channels=input_shape,
+                                        out_channels=hidden_units,
+                                  kernel_size=3,  # how big is the square that's going over the image?
+                                  stride=1,  # default
+                                        padding=1),  # options = "valid" (no padding) or "same" (output has same shape as input) or int for specific number
+                    nn.ReLU(),
+                        nn.Conv2d(in_channels=hidden_units,
+                                  out_channels=hidden_units,
+                                  kernel_size=3,
+                                  stride=1,
+                                  padding=1),
+                    nn.ReLU(),
+                    nn.MaxPool2d(kernel_size=2,
+                                     stride=2)  # default stride value is same as kernel_size
+                )
+                self.block_2 = nn.Sequential(
+                    nn.Conv2d(hidden_units, hidden_units, 3, padding=1),
+                    nn.ReLU(),
+                    nn.Conv2d(hidden_units, hidden_units, 3, padding=1),
+                    nn.ReLU(),
+                    nn.MaxPool2d(2)
+                )
+                self.classifier = nn.Sequential(
+                    nn.Flatten(),
+                        nn.Linear(in_features=hidden_units*7*7,
+                                  out_features=output_shape)
+                )
 
-    # Change classifier head with random seed for reproducibility
-    torch.manual_seed(seed)
-    model.heads = nn.Sequential(nn.Linear(in_features=768, # keep this the same as original model
-                                          out_features=num_classes)) # update to reflect target number of classes
-    
-    return model, transforms
+            def forward(self, x: torch.Tensor):
+                # x = self.block_1(x)
+                # print(x.shape)
+                # x = self.block_2(x)
+                # print(x.shape)
+                # x = self.classifier(x)
+                # print(x.shape)
+                x = self.classifier(self.block_2(self.block_1(x)))
+                return x
+    return Model, transforms

model_food101_20_percent.pth → model_1.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b4357a334ed5737baacaf7a99b0ba491ef88c61580790277acec9ef877cd77c9
-size 343564561
+oid sha256:09d5f31bc58b2ae0b7b58d3730491a2708db1245fbaf688d1d6a3cb1b613ba3d
+size 77575