helper_function.py

import torch 
from torch import nn

from timeit import default_timer as timer 
from tqdm.auto import tqdm

import numpy as np
import matplotlib.pyplot as plt

#------------------------------------------------------------------------------
# Accuracy metric: Classification (single or multi-class)
#------------------------------------------------------------------------------
def accyracy_fn(y_true, y_pred):
    """Calculates the accuracy of the predictions

    Args:
        y_true (torch.tensor): tensor values of the true labels
        y_pred (torch.tensor): tensor values of the predicted labels from the model
    
    Returns:
        acc (float): The calculated accuracy of the model predictions (y_pred)
    """
    correct = torch.eq(y_true, y_pred).sum().item() # total number of correct predictions
    acc     = ( correct/len(y_pred) )*100           # len(y_pred) 
    return acc

#------------------------------------------------------------------------------
# Training time: Time taken to run the code
#------------------------------------------------------------------------------
def print_train_time(start: float, end: float, device):
    """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"Train time on {device}: {total_time:.3f} seconds")
    return total_time

#------------------------------------------------------------------------------
# Model evaluation function: Only with model created with a class 
#------------------------------------------------------------------------------
def eval_mode(model: torch.nn.Module,
              data_loader: torch.utils.data.DataLoader,
              loss_fn: torch.nn.Module,
              accyracy_fn,
              device: torch.device):
    """Reutrns a dictionary containig the results of model perdicting on data_loader"""
    loss, acc = 0, 0
    model.eval()
    with torch.inference_mode():
        for X_batch,y_batch in tqdm(data_loader):
            # Put data on target device
            X_batch, y_batch = X_batch.to(device), y_batch.to(device)
            # Make predictions
            y_pred = model(X_batch)

            # Accumulate the loss and values per batch
            loss += loss_fn(y_pred, y_batch)
            acc += accyracy_fn(y_batch, y_pred.argmax(dim=1))

        
        # Scale loss & acc to find the average loss/acc per bath
        avgloss = loss/len(data_loader)
        avgacc  = acc/len(data_loader)

        return {"model_name": model.__class__.__name__, # only works when model was created with a class
                "model_loss": avgloss.item(),
                "model_acc": avgacc}
    
#------------------------------------------------------------------------------
# Model training function:
#------------------------------------------------------------------------------
def train_step(model: torch.nn.Module,
                data_loader: torch.utils.data.DataLoader,
                loss_fn: torch.nn.Module,
                optimizer: torch.optim.Optimizer,
                accuracy_fn,
                device: torch.device):
    """Performs a training with model learning on data_loader"""
    train_loss, train_acc = 0, 0

    # Batch pass: loop through the training batches
    for batch, (X_batch, y_batch) in enumerate(data_loader):
        # Put data on target device
        X_batch, y_batch = X_batch.to(device), y_batch.to(device)

        # 1. Forward propogation of batch
        y_pred = model(X_batch)

        # 2. Calculate loss & accuracy (per batch)
        loss_temp   = loss_fn(y_pred, y_batch)
        train_loss  += loss_temp  # Accumulate train loss

        train_acc   += accuracy_fn(y_true=y_batch,
                                   y_pred=y_pred.argmax(dim=1)) # argmax to go from logits -> prediction labes
        
        # 3. Optimizer zero grad
        optimizer.zero_grad()

        # 4. Loss backward
        loss_temp.backward()

        # 5. Optimizer step (update model parameters once per batch)
        optimizer.step()

        # Verbose **per batch**
        if batch % 400 == 0:
            print(f"Batch number: {batch} | Samples trained over: { batch * len(X_batch) }/{ len(X_batch) }")

    # Adjust metrics
    train_loss  /= len(data_loader) # Averaged (over batch) train loss
    train_acc   /= len(data_loader) # Averaged (over batch) train accuracy

    # Verbose **over all the batches**
    print(f"Train loss: {train_loss:.3f} | Train acc: {train_acc} %")


def test_step(model: torch.nn.Module,
              data_loader: torch.utils.data.DataLoader,
              loss_fn: torch.nn.Module,
              accuracy_fn,
              device: torch.device):
    """Perform a testing loop step on model going over data_loader"""
    test_loss, test_acc = 0, 0

    # Set the model to eval mode
    model.eval()

    # Turn on the inference mode context manager
    with torch.inference_mode():
        for X_batch, y_batch in data_loader:
            # Send the data to the target device
            X_batch, y_batch = X_batch.to(device), y_batch.to(device)

            # 1. Forward propogation of batch
            test_pred = model(X_batch)

            # 2. Calculate the loss & accuracy of model for test data **per batch**
            test_loss   += loss_fn(test_pred,y_batch)
            test_acc    += accuracy_fn(y_true=y_batch,
                                       y_pred=test_pred.argmax(dim=1)) # go from logits -> prediction labels
            
        
        # Adjust mertics
        test_loss   /= len(data_loader) # Averaged (over batch) test loss
        test_acc    /= len(data_loader) # Averaged (over batch) test accuracy

        # Verbose **over all the batches**
        print(f"Test loss: {test_loss:.3f} | Test acc: {test_acc:.3f} %\n")