Optimised the files

models/custom_resnet.py

@@ -0,0 +1,456 @@
+"""Module to define the model."""
+
+# Resources
+# https://lightning.ai/docs/pytorch/stable/starter/introduction.html
+# https://lightning.ai/docs/pytorch/stable/starter/converting.html
+# https://lightning.ai/docs/pytorch/stable/notebooks/lightning_examples/cifar10-baseline.html
+
+import modules.config as config
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torchinfo
+from torch.optim.lr_scheduler import OneCycleLR
+from torch_lr_finder import LRFinder
+from torchmetrics import Accuracy
+
+# What is the start LR and weight decay you'd prefer?
+PREFERRED_START_LR = config.PREFERRED_START_LR
+PREFERRED_WEIGHT_DECAY = config.PREFERRED_WEIGHT_DECAY
+
+
+def detailed_model_summary(model, input_size):
+    """Define a function to print the model summary."""
+
+    # https://github.com/TylerYep/torchinfo
+    torchinfo.summary(
+        model,
+        input_size=input_size,
+        batch_dim=0,
+        col_names=(
+            "input_size",
+            "kernel_size",
+            "output_size",
+            "num_params",
+            "trainable",
+        ),
+        verbose=1,
+        col_width=16,
+    )
+
+
+############# Assignment 13 Model #############
+
+
+# This is for Assignment 13
+# Model used from Assignment 11 and converted to lightning model
+class CustomResNet(pl.LightningModule):
+    """This defines the structure of the NN."""
+
+    # Class variable to print shape
+    print_shape = False
+    # Default dropout value
+    dropout_value = 0.02
+
+    def __init__(self):
+        super().__init__()
+
+        # Define loss function
+        # https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html
+        self.loss_function = torch.nn.CrossEntropyLoss()
+
+        # Define accuracy function
+        # https://torchmetrics.readthedocs.io/en/stable/classification/accuracy.html
+        self.accuracy_function = Accuracy(task="multiclass", num_classes=10)
+
+        # Add results dictionary
+        self.results = {
+            "train_loss": [],
+            "train_acc": [],
+            "test_loss": [],
+            "test_acc": [],
+            "val_loss": [],
+            "val_acc": [],
+        }
+
+        # Save misclassified images
+        self.misclassified_image_data = {"images": [], "ground_truths": [], "predicted_vals": []}
+
+        # LR
+        self.learning_rate = PREFERRED_START_LR
+
+        #  Model Notes
+
+        # PrepLayer - Conv 3x3 s1, p1) >> BN >> RELU [64k]
+        # 1. Input size: 32x32x3
+        self.prep = nn.Sequential(
+            nn.Conv2d(
+                in_channels=3,
+                out_channels=64,
+                kernel_size=(3, 3),
+                stride=1,
+                padding=1,
+                dilation=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.Dropout(self.dropout_value),
+        )
+
+        # Layer1: X = Conv 3x3 (s1, p1) >> MaxPool2D >> BN >> RELU [128k]
+        self.layer1_x = nn.Sequential(
+            nn.Conv2d(
+                in_channels=64,
+                out_channels=128,
+                kernel_size=(3, 3),
+                stride=1,
+                padding=1,
+                dilation=1,
+                bias=False,
+            ),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            nn.Dropout(self.dropout_value),
+        )
+
+        # Layer1: R1 = ResBlock( (Conv-BN-ReLU-Conv-BN-ReLU))(X) [128k]
+        self.layer1_r1 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=128,
+                out_channels=128,
+                kernel_size=(3, 3),
+                stride=1,
+                padding=1,
+                dilation=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            nn.Dropout(self.dropout_value),
+            nn.Conv2d(
+                in_channels=128,
+                out_channels=128,
+                kernel_size=(3, 3),
+                stride=1,
+                padding=1,
+                dilation=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            nn.Dropout(self.dropout_value),
+        )
+
+        # Layer 2: Conv 3x3 [256k], MaxPooling2D, BN, ReLU
+        self.layer2 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=128,
+                out_channels=256,
+                kernel_size=(3, 3),
+                stride=1,
+                padding=1,
+                dilation=1,
+                bias=False,
+            ),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.BatchNorm2d(256),
+            nn.ReLU(),
+            nn.Dropout(self.dropout_value),
+        )
+
+        # Layer 3: X = Conv 3x3 (s1, p1) >> MaxPool2D >> BN >> RELU [512k]
+        self.layer3_x = nn.Sequential(
+            nn.Conv2d(
+                in_channels=256,
+                out_channels=512,
+                kernel_size=(3, 3),
+                stride=1,
+                padding=1,
+                dilation=1,
+                bias=False,
+            ),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.BatchNorm2d(512),
+            nn.ReLU(),
+            nn.Dropout(self.dropout_value),
+        )
+
+        # Layer 3: R2 = ResBlock( (Conv-BN-ReLU-Conv-BN-ReLU))(X) [512k]
+        self.layer3_r2 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=512,
+                out_channels=512,
+                kernel_size=(3, 3),
+                stride=1,
+                padding=1,
+                dilation=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(512),
+            nn.ReLU(),
+            nn.Dropout(self.dropout_value),
+            nn.Conv2d(
+                in_channels=512,
+                out_channels=512,
+                kernel_size=(3, 3),
+                stride=1,
+                padding=1,
+                dilation=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(512),
+            nn.ReLU(),
+            nn.Dropout(self.dropout_value),
+        )
+
+        # MaxPooling with Kernel Size 4
+        # If stride is None, it is set to kernel_size
+        self.maxpool = nn.MaxPool2d(kernel_size=4, stride=4)
+
+        # FC Layer
+        self.fc = nn.Linear(512, 10)
+
+        # Save hyperparameters
+        self.save_hyperparameters()
+
+    def print_view(self, x, msg=""):
+        """Print shape of the model"""
+        if self.print_shape:
+            if msg != "":
+                print(msg, "\n\t", x.shape, "\n")
+            else:
+                print(x.shape)
+
+    def forward(self, x):
+        """Forward pass"""
+
+        # PrepLayer
+        x = self.prep(x)
+        self.print_view(x, "PrepLayer")
+
+        # Layer 1
+        x = self.layer1_x(x)
+        self.print_view(x, "Layer 1, X")
+        r1 = self.layer1_r1(x)
+        self.print_view(r1, "Layer 1, R1")
+        x = x + r1
+        self.print_view(x, "Layer 1, X + R1")
+
+        # Layer 2
+        x = self.layer2(x)
+        self.print_view(x, "Layer 2")
+
+        # Layer 3
+        x = self.layer3_x(x)
+        self.print_view(x, "Layer 3, X")
+        r2 = self.layer3_r2(x)
+        self.print_view(r2, "Layer 3, R2")
+        x = x + r2
+        self.print_view(x, "Layer 3, X + R2")
+
+        # MaxPooling
+        x = self.maxpool(x)
+        self.print_view(x, "Max Pooling")
+
+        # FC Layer
+        # Reshape before FC such that it becomes 1D
+        x = x.view(x.shape[0], -1)
+        self.print_view(x, "Reshape before FC")
+        x = self.fc(x)
+        self.print_view(x, "After FC")
+
+        # Softmax
+        return F.log_softmax(x, dim=-1)
+
+    # Alert: Remove this function later as Tuner is now being used to automatically find the best LR
+    def find_optimal_lr(self, train_loader):
+        """Use LR Finder to find the best starting learning rate"""
+
+        # https://github.com/davidtvs/pytorch-lr-finder
+        # https://github.com/davidtvs/pytorch-lr-finder#notes
+        # https://github.com/davidtvs/pytorch-lr-finder/blob/master/torch_lr_finder/lr_finder.py
+
+        # New optimizer with default LR
+        tmp_optimizer = optim.Adam(self.parameters(), lr=PREFERRED_START_LR, weight_decay=PREFERRED_WEIGHT_DECAY)
+
+        # Create LR finder object
+        lr_finder = LRFinder(self, optimizer=tmp_optimizer, criterion=self.loss_function)
+        lr_finder.range_test(train_loader=train_loader, end_lr=10, num_iter=100)
+        # https://github.com/davidtvs/pytorch-lr-finder/issues/88
+        _, suggested_lr = lr_finder.plot(suggest_lr=True)
+        lr_finder.reset()
+        # plot.figure.savefig("LRFinder - Suggested Max LR.png")
+
+        print(f"Suggested Max LR: {suggested_lr}")
+
+        if suggested_lr is None:
+            suggested_lr = PREFERRED_START_LR
+
+        return suggested_lr
+
+    # optimiser function
+    def configure_optimizers(self):
+        """Add ADAM optimizer to the lightning module"""
+        optimizer = optim.Adam(self.parameters(), lr=self.learning_rate, weight_decay=PREFERRED_WEIGHT_DECAY)
+
+        # Percent start for OneCycleLR
+        # Handles the case where max_epochs is less than 5
+        percent_start = 5 / int(self.trainer.max_epochs)
+        if percent_start >= 1:
+            percent_start = 0.3
+
+        # https://lightning.ai/docs/pytorch/stable/common/optimization.html#total-stepping-batches
+        scheduler_dict = {
+            "scheduler": OneCycleLR(
+                optimizer=optimizer,
+                max_lr=self.learning_rate,
+                total_steps=int(self.trainer.estimated_stepping_batches),
+                pct_start=percent_start,
+                div_factor=100,
+                three_phase=False,
+                anneal_strategy="linear",
+                final_div_factor=100,
+                verbose=False,
+            ),
+            "interval": "step",
+        }
+
+        return {"optimizer": optimizer, "lr_scheduler": scheduler_dict}
+
+    # Define loss function
+    def compute_loss(self, prediction, target):
+        """Compute Loss"""
+
+        # Calculate loss
+        loss = self.loss_function(prediction, target)
+
+        return loss
+
+    # Define accuracy function
+    def compute_accuracy(self, prediction, target):
+        """Compute accuracy"""
+
+        # Calculate accuracy
+        acc = self.accuracy_function(prediction, target)
+
+        return acc * 100
+
+    # Function to compute loss and accuracy for both training and validation
+    def compute_metrics(self, batch):
+        """Function to calculate loss and accuracy"""
+
+        # Get data and target from batch
+        data, target = batch
+
+        # Generate predictions using model
+        pred = self(data)
+
+        # Calculate loss for the batch
+        loss = self.compute_loss(prediction=pred, target=target)
+
+        # Calculate accuracy for the batch
+        acc = self.compute_accuracy(prediction=pred, target=target)
+
+        return loss, acc
+
+    # Get misclassified images based on how many images to return
+    def store_misclassified_images(self):
+        """Get an array of misclassified images"""
+
+        self.misclassified_image_data = {"images": [], "ground_truths": [], "predicted_vals": []}
+
+        # Initialize the model to evaluation mode
+        self.eval()
+
+        # Disable gradient calculation while testing
+        with torch.no_grad():
+            for batch in self.trainer.test_dataloaders:
+                # Move data and labels to device
+                data, target = batch
+                data, target = data.to(self.device), target.to(self.device)
+
+                # Predict using model
+                pred = self(data)
+
+                # Get the index of the max log-probability
+                output = pred.argmax(dim=1)
+
+                # Save the incorrect predictions
+                incorrect_indices = ~output.eq(target)
+
+                # Store images incorrectly predicted, generated predictions and the actual value
+                self.misclassified_image_data["images"].extend(data[incorrect_indices])
+                self.misclassified_image_data["ground_truths"].extend(target[incorrect_indices])
+                self.misclassified_image_data["predicted_vals"].extend(output[incorrect_indices])
+
+    # training function
+    def training_step(self, batch, batch_idx):
+        """Training step"""
+
+        # Compute loss and accuracy
+        loss, acc = self.compute_metrics(batch)
+
+        self.log("train_loss", loss, prog_bar=True, on_epoch=True, logger=True)
+        self.log("train_acc", acc, prog_bar=True, on_epoch=True, logger=True)
+        # Return training loss
+        return loss
+
+    # validation function
+    def validation_step(self, batch, batch_idx):
+        """Validation step"""
+
+        # Compute loss and accuracy
+        loss, acc = self.compute_metrics(batch)
+
+        self.log("val_loss", loss, prog_bar=True, on_epoch=True, logger=True)
+        self.log("val_acc", acc, prog_bar=True, on_epoch=True, logger=True)
+        # Return validation loss
+        return loss
+
+    # test function will just use validation step
+    def test_step(self, batch, batch_idx):
+        """Test step"""
+
+        # Compute loss and accuracy
+        loss, acc = self.compute_metrics(batch)
+
+        self.log("test_loss", loss, prog_bar=False, on_epoch=True, logger=True)
+        self.log("test_acc", acc, prog_bar=False, on_epoch=True, logger=True)
+        # Return validation loss
+        return loss
+
+    # At the end of train epoch append the training loss and accuracy to an instance variable called results
+    def on_train_epoch_end(self):
+        """On train epoch end"""
+
+        # Append training loss and accuracy to results
+        self.results["train_loss"].append(self.trainer.callback_metrics["train_loss"].detach().item())
+        self.results["train_acc"].append(self.trainer.callback_metrics["train_acc"].detach().item())
+
+    # At the end of validation epoch append the validation loss and accuracy to an instance variable called results
+    def on_validation_epoch_end(self):
+        """On validation epoch end"""
+
+        # Append validation loss and accuracy to results
+        self.results["test_loss"].append(self.trainer.callback_metrics["val_loss"].detach().item())
+        self.results["test_acc"].append(self.trainer.callback_metrics["val_acc"].detach().item())
+
+    # # At the end of test epoch append the test loss and accuracy to an instance variable called results
+    # def on_test_epoch_end(self):
+    #     """On test epoch end"""
+
+    #     # Append test loss and accuracy to results
+    #     self.results["test_loss"].append(self.trainer.callback_metrics["test_loss"].detach().item())
+    #     self.results["test_acc"].append(self.trainer.callback_metrics["test_acc"].detach().item())
+
+    # At the end of test save misclassified images, the predictions and ground truth in an instance variable called misclassified_image_data
+    def on_test_end(self):
+        """On test end"""
+
+        print("Test ended! Saving misclassified images")
+        # Get misclassified images
+        self.store_misclassified_images()

modules/config.py

@@ -0,0 +1,50 @@
+# Alert: Change these when running in production
+
+# Constants naming convention: All caps separated by underscore
+# https://realpython.com/python-constants/
+
+# Where do we store the data?
+DATA_PATH = "../../data/"
+CHECKPOINT_PATH = "../../checkpoints/"
+LOGGING_PATH = "../../logs/"
+MISCLASSIFIED_PATH = "Misclassified_Data.pt"
+MODEL_PATH = "CustomResNet.pt"
+
+# Specify the number of epochs
+NUM_EPOCHS = 24
+
+# Set the batch size
+BATCH_SIZE = 512
+
+# Set seed value for reproducibility
+SEED = 53
+
+# What is the start LR and weight decay you'd prefer?
+PREFERRED_START_LR = 5e-3
+PREFERRED_WEIGHT_DECAY = 1e-5
+
+
+# What is the mean and std deviation of the dataset?
+CIFAR_MEAN = (0.4915, 0.4823, 0.4468)
+CIFAR_STD = (0.2470, 0.2435, 0.2616)
+
+# What is the cutout size?
+CUTOUT_SIZE = 16
+
+# What are the classes in CIFAR10?
+# Create class labels and convert to tuple
+CIFAR_CLASSES = tuple(
+    c.capitalize()
+    for c in [
+        "plane",
+        "car",
+        "bird",
+        "cat",
+        "deer",
+        "dog",
+        "frog",
+        "horse",
+        "ship",
+        "truck",
+    ]
+)

modules/dataset.py

@@ -0,0 +1,110 @@
+"""This file contains functions to download and transform the CIFAR10 dataset"""
+# Needed for image transformations
+import albumentations as A
+import modules.config as config
+
+# # Needed for padding issues in albumentations
+# import cv2
+import numpy as np
+from albumentations.pytorch.transforms import ToTensorV2
+from torch.utils.data import Dataset
+
+# Use precomputed values for mean and standard deviation of the dataset
+CIFAR_MEAN = config.CIFAR_MEAN
+CIFAR_STD = config.CIFAR_STD
+CUTOUT_SIZE = config.CUTOUT_SIZE
+
+# Create class labels and convert to tuple
+CIFAR_CLASSES = config.CIFAR_CLASSES
+
+
+class CIFAR10Transforms(Dataset):
+    """Apply albumentations augmentations to CIFAR10 dataset"""
+
+    # Given a dataset and transformations,
+    # apply the transformations and return the dataset
+    def __init__(self, dataset, transforms):
+        self.dataset = dataset
+        self.transforms = transforms
+
+    def __getitem__(self, idx):
+        # Get the image and label from the dataset
+        image, label = self.dataset[idx]
+
+        # Apply transformations on the image
+        image = self.transforms(image=np.array(image))["image"]
+
+        return image, label
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __repr__(self):
+        return f"CIFAR10Transforms(dataset={self.dataset}, transforms={self.transforms})"
+
+    def __str__(self):
+        return f"CIFAR10Transforms(dataset={self.dataset}, transforms={self.transforms})"
+
+
+def apply_cifar_image_transformations(mean=CIFAR_MEAN, std=CIFAR_STD, cutout_size=CUTOUT_SIZE):
+    """
+    Function to apply the required transformations to the MNIST dataset.
+    """
+    # Apply the required transformations to the MNIST dataset
+    train_transforms = A.Compose(
+        [
+            # normalize the images with mean and standard deviation from the whole dataset
+            # https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Normalize
+            # # transforms.Normalize(cifar_mean, cifar_std),
+            A.Normalize(mean=list(mean), std=list(std)),
+            # RandomCrop 32, 32 (after padding of 4)
+            # https://albumentations.ai/docs/api_reference/augmentations/geometric/transforms/#albumentations.augmentations.geometric.transforms.PadIfNeeded
+            # MinHeight and MinWidth are set to 36 to ensure that the image is padded to 36x36 after padding
+            # border_mode (OpenCV flag): flag that is used to specify the pixel extrapolation method. Should be one of:
+            # cv2.BORDER_CONSTANT, cv2.BORDER_REPLICATE, cv2.BORDER_REFLECT, cv2.BORDER_WRAP, cv2.BORDER_REFLECT_101.
+            # Default: cv2.BORDER_REFLECT_101
+            A.PadIfNeeded(min_height=36, min_width=36),
+            # https://albumentations.ai/docs/api_reference/augmentations/crops/transforms/#albumentations.augmentations.crops.transforms.RandomCrop
+            A.RandomCrop(32, 32),
+            # CutOut(8, 8)
+            # # https://albumentations.ai/docs/api_reference/augmentations/dropout/cutout/#albumentations.augmentations.dropout.cutout.Cutout
+            # # Because we normalized the images with mean and standard deviation from the whole dataset, the fill_value is set to the mean of the dataset
+            # A.Cutout(
+            #     num_holes=1, max_h_size=cutout_size, max_w_size=cutout_size, p=1.0
+            # ),
+            # https://albumentations.ai/docs/api_reference/augmentations/dropout/coarse_dropout/#coarsedropout-augmentation-augmentationsdropoutcoarse_dropout
+            A.CoarseDropout(
+                max_holes=1,
+                max_height=cutout_size,
+                max_width=cutout_size,
+                min_holes=1,
+                min_height=cutout_size,
+                min_width=cutout_size,
+                p=1.0,
+            ),
+            # Convert the images to tensors
+            # # transforms.ToTensor(),
+            ToTensorV2(),
+        ]
+    )
+
+    # Test data transformations
+    test_transforms = A.Compose(
+        # Convert the images to tensors
+        # normalize the images with mean and standard deviation from the whole dataset
+        [
+            A.Normalize(mean=list(mean), std=list(std)),
+            # Convert the images to tensors
+            ToTensorV2(),
+        ]
+    )
+
+    return train_transforms, test_transforms
+
+
+def calculate_mean_std(dataset):
+    """Function to calculate the mean and standard deviation of CIFAR dataset"""
+    data = dataset.data.astype(np.float32) / 255.0
+    mean = np.mean(data, axis=(0, 1, 2))
+    std = np.std(data, axis=(0, 1, 2))
+    return mean, std

modules/lightning_dataset.py

@@ -0,0 +1,109 @@
+"""This file contains functions to prepare dataloader in the way lightning expects"""
+import pytorch_lightning as pl
+import torchvision.datasets as datasets
+from lightning_fabric.utilities.seed import seed_everything
+from modules.dataset import CIFAR10Transforms, apply_cifar_image_transformations
+from torch.utils.data import DataLoader, random_split
+
+
+class CIFARDataModule(pl.LightningDataModule):
+    """Lightning DataModule for CIFAR10 dataset"""
+
+    def __init__(self, data_path, batch_size, seed, val_split=0, num_workers=0):
+        super().__init__()
+
+        self.data_path = data_path
+        self.batch_size = batch_size
+        self.seed = seed
+        self.val_split = val_split
+        self.num_workers = num_workers
+        self.dataloader_dict = {
+            # "shuffle": True,
+            "batch_size": self.batch_size,
+            "num_workers": self.num_workers,
+            "pin_memory": True,
+            # "worker_init_fn": self._init_fn,
+            "persistent_workers": self.num_workers > 0,
+        }
+        self.prepare_data_per_node = False
+
+        # Fixes attribute defined outside __init__ warning
+        self.training_dataset = None
+        self.validation_dataset = None
+        self.testing_dataset = None
+
+        # # Make sure data is downloaded
+        # self.prepare_data()
+
+    def _split_train_val(self, dataset):
+        """Split the dataset into train and validation sets"""
+
+        # Throw an error if the validation split is not between 0 and 1
+        if not 0 < self.val_split < 1:
+            raise ValueError("Validation split must be between 0 and 1")
+
+        # # Set seed again, might not be necessary
+        # seed_everything(int(self.seed))
+
+        # Calculate lengths of each dataset
+        total_length = len(dataset)
+        train_length = int((1 - self.val_split) * total_length)
+        val_length = total_length - train_length
+
+        # Split the dataset
+        train_dataset, val_dataset = random_split(dataset, [train_length, val_length])
+
+        return train_dataset, val_dataset
+
+    # https://lightning.ai/docs/pytorch/stable/data/datamodule.html#prepare-data
+    def prepare_data(self):
+        # Download the CIFAR10 dataset if it doesn't exist
+        datasets.CIFAR10(self.data_path, train=True, download=True)
+        datasets.CIFAR10(self.data_path, train=False, download=True)
+
+    # https://lightning.ai/docs/pytorch/stable/data/datamodule.html#setup
+    # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.core.hooks.DataHooks.html#lightning.pytorch.core.hooks.DataHooks.setup
+    def setup(self, stage=None):
+        # seed_everything(int(self.seed))
+
+        # Define the data transformations
+        train_transforms, test_transforms = apply_cifar_image_transformations()
+        val_transforms = test_transforms
+
+        # Create train and validation datasets
+        if stage == "fit" or stage is None:
+            if self.val_split != 0:
+                # Split the training data into training and validation sets
+                data_train, data_val = self._split_train_val(datasets.CIFAR10(self.data_path, train=True))
+                # Apply transformations
+                self.training_dataset = CIFAR10Transforms(data_train, train_transforms)
+                self.validation_dataset = CIFAR10Transforms(data_val, val_transforms)
+            else:
+                # Only training data here
+                self.training_dataset = CIFAR10Transforms(
+                    datasets.CIFAR10(self.data_path, train=True), train_transforms
+                )
+                # Validation will be same sa test
+                self.validation_dataset = CIFAR10Transforms(
+                    datasets.CIFAR10(self.data_path, train=False), val_transforms
+                )
+
+        # Create test dataset
+        if stage == "test" or stage is None:
+            # Assign Test split(s) for use in Dataloaders
+            self.testing_dataset = CIFAR10Transforms(datasets.CIFAR10(self.data_path, train=False), test_transforms)
+
+    # https://lightning.ai/docs/pytorch/stable/data/datamodule.html#train-dataloader
+    def train_dataloader(self):
+        return DataLoader(self.training_dataset, **self.dataloader_dict, shuffle=True)
+
+    # https://lightning.ai/docs/pytorch/stable/data/datamodule.html#val-dataloader
+    def val_dataloader(self):
+        return DataLoader(self.validation_dataset, **self.dataloader_dict, shuffle=False)
+
+    # https://lightning.ai/docs/pytorch/stable/data/datamodule.html#test-dataloader
+    def test_dataloader(self):
+        return DataLoader(self.testing_dataset, **self.dataloader_dict, shuffle=False)
+
+    def _init_fn(self, worker_id):
+        seed_everything(int(self.seed) + worker_id)

modules/trainer.py

@@ -0,0 +1,120 @@
+"""Module to define the train and test functions."""
+
+# from functools import partial
+
+import modules.config as config
+import pytorch_lightning as pl
+import torch
+from modules.utils import create_folder_if_not_exists
+from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint, ModelSummary
+
+# Import tuner
+from pytorch_lightning.tuner.tuning import Tuner
+
+# What is the start LR and weight decay you'd prefer?
+PREFERRED_START_LR = config.PREFERRED_START_LR
+
+
+def train_and_test_model(
+    batch_size,
+    num_epochs,
+    model,
+    datamodule,
+    logger,
+    debug=False,
+):
+    """Trains and tests the model by iterating through epochs using Lightning Trainer."""
+
+    print(f"\n\nBatch size: {batch_size}, Total epochs: {num_epochs}\n\n")
+
+    print("Defining Lightning Callbacks")
+
+    # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html#modelcheckpoint
+    checkpoint = ModelCheckpoint(
+        dirpath=config.CHECKPOINT_PATH, monitor="val_acc", mode="max", filename="model_best_epoch", save_last=True
+    )
+    # # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.LearningRateMonitor.html#learningratemonitor
+    lr_rate_monitor = LearningRateMonitor(logging_interval="epoch", log_momentum=False)
+    # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelSummary.html#lightning.pytorch.callbacks.ModelSummary
+    model_summary = ModelSummary(max_depth=0)
+
+    print("Defining Lightning Trainer")
+    # Change trainer settings for debugging
+    if debug:
+        num_epochs = 1
+        fast_dev_run = True
+        overfit_batches = 0.1
+        profiler = "advanced"
+    else:
+        fast_dev_run = False
+        overfit_batches = 0.0
+        profiler = None
+
+    # https://lightning.ai/docs/pytorch/stable/common/trainer.html#methods
+    trainer = pl.Trainer(
+        precision=16,
+        fast_dev_run=fast_dev_run,
+        # deterministic=True,
+        # devices="auto",
+        # accelerator="auto",
+        max_epochs=num_epochs,
+        logger=logger,
+        # enable_model_summary=False,
+        overfit_batches=overfit_batches,
+        log_every_n_steps=10,
+        # num_sanity_val_steps=5,
+        profiler=profiler,
+        # check_val_every_n_epoch=1,
+        callbacks=[checkpoint, lr_rate_monitor, model_summary],
+        # callbacks=[checkpoint],
+    )
+
+    # # Using the learning rate finder
+    # model.learning_rate = model.find_optimal_lr(train_loader=datamodule.train_dataloader())
+
+    # Using the lr_find from Trainer.tune method instead
+    # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.tuner.tuning.Tuner.html#lightning.pytorch.tuner.tuning.Tuner
+    # https://www.youtube.com/watch?v=cLZv0eZQSIE
+    print("Finding the optimal learning rate using Lightning Tuner.")
+    tuner = Tuner(trainer)
+    tuner.lr_find(
+        model=model,
+        datamodule=datamodule,
+        min_lr=PREFERRED_START_LR,
+        max_lr=5,
+        num_training=200,
+        mode="linear",
+        early_stop_threshold=10,
+        attr_name="learning_rate",
+    )
+
+    trainer.fit(model, datamodule=datamodule)
+    trainer.test(model, dataloaders=datamodule.test_dataloader())
+
+    # # Obtain the results dictionary from model
+    print("Collecting epoch level model results.")
+    results = model.results
+    # print(f"Results Length: {len(results)}")
+
+    # Get the list of misclassified images
+    print("Collecting misclassified images.")
+    misclassified_image_data = model.misclassified_image_data
+    # print(f"Misclassified Images Length: {len(misclassified_image_data)}")
+
+    # Save the model using torch save as backup
+    print("Saving the model.")
+    create_folder_if_not_exists(config.MODEL_PATH)
+    torch.save(model.state_dict(), config.MODEL_PATH)
+
+    # Save first few misclassified images data to a file
+    num_elements = 20
+    print(f"Saving first {num_elements} misclassified images.")
+    subset_misclassified_image_data = {"images": [], "ground_truths": [], "predicted_vals": []}
+    subset_misclassified_image_data["images"] = misclassified_image_data["images"][:num_elements]
+    subset_misclassified_image_data["ground_truths"] = misclassified_image_data["ground_truths"][:num_elements]
+    subset_misclassified_image_data["predicted_vals"] = misclassified_image_data["predicted_vals"][:num_elements]
+    create_folder_if_not_exists(config.MISCLASSIFIED_PATH)
+    torch.save(subset_misclassified_image_data, config.MISCLASSIFIED_PATH)
+
+    return trainer, results, misclassified_image_data
+    # return trainer

modules/utils.py

@@ -0,0 +1,70 @@
+"""Module to define utility functions for the project."""
+import os
+
+import torch
+
+
+def get_num_workers(model_run_location):
+    """Given a run mode, return the number of workers to be used for data loading."""
+
+    # calculate the number of workers
+    num_workers = (os.cpu_count() - 1) if os.cpu_count() > 3 else 2
+
+    # If run_mode is local, use only 2 workers
+    num_workers = num_workers if model_run_location == "colab" else 0
+
+    return num_workers
+
+
+# Function to save the model
+# https://debuggercafe.com/saving-and-loading-the-best-model-in-pytorch/
+def save_model(epoch, model, optimizer, scheduler, batch_size, criterion, file_name):
+    """
+    Function to save the trained model along with other information to disk.
+    """
+    # print(f"Saving model from epoch {epoch}...")
+    torch.save(
+        {
+            "epoch": epoch,
+            "model_state_dict": model.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "scheduler_state_dict": scheduler.state_dict(),
+            "batch_size": batch_size,
+            "loss": criterion,
+        },
+        file_name,
+    )
+
+
+# Given a list of train_losses, train_accuracies, test_losses,
+# test_accuracies, loop through epoch and print the metrics
+def pretty_print_metrics(num_epochs, results):
+    """
+    Function to print the metrics in a pretty format.
+    """
+    # Extract train_losses, train_acc, test_losses, test_acc from results
+    train_losses = results["train_loss"]
+    train_acc = results["train_acc"]
+    test_losses = results["test_loss"]
+    test_acc = results["test_acc"]
+
+    for i in range(num_epochs):
+        print(
+            f"Epoch: {i+1:02d}, Train Loss: {train_losses[i]:.4f}, "
+            f"Test Loss: {test_losses[i]:.4f}, Train Accuracy: {train_acc[i]:.4f}, "
+            f"Test Accuracy: {test_acc[i]:.4f}"
+        )
+
+
+# Given a file path, extract the folder path and create folder recursively if it does not already exist
+def create_folder_if_not_exists(file_path):
+    """
+    Function to create a folder if it does not exist.
+    """
+    # Extract the folder path
+    folder_path = os.path.dirname(file_path)
+
+    # Create the folder if it does not exist
+    if not os.path.exists(folder_path):
+        os.makedirs(folder_path)
+        print(f"Created folder: {folder_path}")

modules/visualize.py

@@ -0,0 +1,169 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+
+
+def convert_back_image(image):
+    """Using mean and std deviation convert image back to normal"""
+    cifar10_mean = (0.4914, 0.4822, 0.4471)
+    cifar10_std = (0.2469, 0.2433, 0.2615)
+    image = image.numpy().astype(dtype=np.float32)
+
+    for i in range(image.shape[0]):
+        image[i] = (image[i] * cifar10_std[i]) + cifar10_mean[i]
+
+    # To stop throwing a warning that image pixels exceeds bounds
+    image = image.clip(0, 1)
+
+    return np.transpose(image, (1, 2, 0))
+
+
+def plot_sample_training_images(batch_data, batch_label, class_label, num_images=30):
+    """Function to plot sample images from the training data."""
+    images, labels = batch_data, batch_label
+
+    # Calculate the number of images to plot
+    num_images = min(num_images, len(images))
+    # calculate the number of rows and columns to plot
+    num_cols = 5
+    num_rows = int(np.ceil(num_images / num_cols))
+
+    # Initialize a subplot with the required number of rows and columns
+    fig, axs = plt.subplots(num_rows, num_cols, figsize=(10, 10))
+
+    # Iterate through the images and plot them in the grid along with class labels
+
+    for img_index in range(1, num_images + 1):
+        plt.subplot(num_rows, num_cols, img_index)
+        plt.tight_layout()
+        plt.axis("off")
+        plt.imshow(convert_back_image(images[img_index - 1]))
+        plt.title(class_label[labels[img_index - 1].item()])
+        plt.xticks([])
+        plt.yticks([])
+
+    return fig, axs
+
+
+def plot_train_test_metrics(results):
+    """
+    Function to plot the training and test metrics.
+    """
+    # Extract train_losses, train_acc, test_losses, test_acc from results
+    train_losses = results["train_loss"]
+    train_acc = results["train_acc"]
+    test_losses = results["test_loss"]
+    test_acc = results["test_acc"]
+
+    # Plot the graphs in a 1x2 grid showing the training and test metrics
+    fig, axs = plt.subplots(1, 2, figsize=(16, 8))
+
+    # Loss plot
+    axs[0].plot(train_losses, label="Train")
+    axs[0].plot(test_losses, label="Test")
+    axs[0].set_title("Loss")
+    axs[0].legend(loc="upper right")
+
+    # Accuracy plot
+    axs[1].plot(train_acc, label="Train")
+    axs[1].plot(test_acc, label="Test")
+    axs[1].set_title("Accuracy")
+    axs[1].legend(loc="upper right")
+
+    return fig, axs
+
+
+def plot_misclassified_images(data, class_label, num_images=10):
+    """Plot the misclassified images from the test dataset."""
+    # Calculate the number of images to plot
+    num_images = min(num_images, len(data["ground_truths"]))
+    # calculate the number of rows and columns to plot
+    num_cols = 5
+    num_rows = int(np.ceil(num_images / num_cols))
+
+    # Initialize a subplot with the required number of rows and columns
+    fig, axs = plt.subplots(num_rows, num_cols, figsize=(num_cols * 2, num_rows * 2))
+
+    # Iterate through the images and plot them in the grid along with class labels
+
+    for img_index in range(1, num_images + 1):
+        # Get the ground truth and predicted labels for the image
+        label = data["ground_truths"][img_index - 1].cpu().item()
+        pred = data["predicted_vals"][img_index - 1].cpu().item()
+        # Get the image
+        image = data["images"][img_index - 1].cpu()
+        # Plot the image
+        plt.subplot(num_rows, num_cols, img_index)
+        plt.tight_layout()
+        plt.axis("off")
+        plt.imshow(convert_back_image(image))
+        plt.title(f"""ACT: {class_label[label]} \nPRED: {class_label[pred]}""")
+        plt.xticks([])
+        plt.yticks([])
+
+    return fig, axs
+
+
+# Function to plot gradcam for misclassified images using pytorch_grad_cam
+def plot_gradcam_images(
+    model,
+    data,
+    class_label,
+    target_layers,
+    targets=None,
+    num_images=10,
+    image_weight=0.25,
+):
+    """Show gradcam for misclassified images"""
+
+    # Calculate the number of images to plot
+    num_images = min(num_images, len(data["ground_truths"]))
+    # calculate the number of rows and columns to plot
+    num_cols = 5
+    num_rows = int(np.ceil(num_images / num_cols))
+
+    # Initialize a subplot with the required number of rows and columns
+    fig, axs = plt.subplots(num_rows, num_cols, figsize=(num_cols * 2, num_rows * 2))
+
+    # Initialize the GradCAM object
+    # https://github.com/jacobgil/pytorch-grad-cam/blob/master/pytorch_grad_cam/grad_cam.py
+    # https://github.com/jacobgil/pytorch-grad-cam/blob/master/pytorch_grad_cam/base_cam.py
+    cam = GradCAM(model=model, target_layers=target_layers, use_cuda=True)
+
+    # Iterate through the images and plot them in the grid along with class labels
+    for img_index in range(1, num_images + 1):
+        # Extract elements from the data dictionary
+        # Get the ground truth and predicted labels for the image
+        label = data["ground_truths"][img_index - 1].cpu().item()
+        pred = data["predicted_vals"][img_index - 1].cpu().item()
+        # Get the image
+        image = data["images"][img_index - 1].cpu()
+
+        # Get the GradCAM output
+        # https://github.com/jacobgil/pytorch-grad-cam/blob/master/pytorch_grad_cam/utils/model_targets.py
+        grad_cam_output = cam(
+            input_tensor=image.unsqueeze(0),
+            targets=targets,
+            aug_smooth=True,
+            eigen_smooth=True,
+        )
+        grad_cam_output = grad_cam_output[0, :]
+
+        # Overlay gradcam on top of numpy image
+        overlayed_image = show_cam_on_image(
+            convert_back_image(image),
+            grad_cam_output,
+            use_rgb=True,
+            image_weight=image_weight,
+        )
+
+        # Plot the image
+        plt.subplot(num_rows, num_cols, img_index)
+        plt.tight_layout()
+        plt.axis("off")
+        plt.imshow(overlayed_image)
+        plt.title(f"""ACT: {class_label[label]} \nPRED: {class_label[pred]}""")
+        plt.xticks([])
+        plt.yticks([])
+    return fig, axs

utilities/callbacks.py

@@ -1,64 +0,0 @@
-import pytorch_lightning as pl
-from pytorch_lightning.callbacks import Callback
-
-from .visualize import plot_model_training_curves
-
-
-class TrainingEndCallback(Callback):
-    def on_train_end(self, trainer, pl_module):
-        # Perform actions at the end of the entire training process
-        print("Training, validation, and testing completed!")
-
-        logged_metrics = pl_module.log_store
-
-        plot_model_training_curves(
-            train_accs=logged_metrics["train_acc_epoch"],
-            test_accs=logged_metrics["val_acc_epoch"],
-            train_losses=logged_metrics["train_loss_epoch"],
-            test_losses=logged_metrics["val_loss_epoch"],
-        )
-
-
-class PrintLearningMetricsCallback(Callback):
-    def on_train_epoch_end(
-        self, trainer: pl.Trainer, pl_module: pl.LightningModule
-    ) -> None:
-        super().on_train_epoch_end(trainer, pl_module)
-        print(
-            f"\nEpoch: {trainer.current_epoch}, Train Loss: {trainer.logged_metrics['train_loss_epoch']}, Train Accuracy: {trainer.logged_metrics['train_acc_epoch']}"
-        )
-        pl_module.log_store.get("train_loss_epoch").append(
-            trainer.logged_metrics["train_loss_epoch"].cpu().detach().item()
-        )
-        pl_module.log_store.get("train_acc_epoch").append(
-            trainer.logged_metrics["train_acc_epoch"].cpu().detach().item()
-        )
-
-    def on_validation_epoch_end(
-        self, trainer: pl.Trainer, pl_module: pl.LightningModule
-    ) -> None:
-        super().on_validation_epoch_end(trainer, pl_module)
-        print(
-            f"\nEpoch: {trainer.current_epoch}, Val Loss: {trainer.logged_metrics['val_loss_epoch']}, Val Accuracy: {trainer.logged_metrics['val_acc_epoch']}"
-        )
-        pl_module.log_store.get("val_loss_epoch").append(
-            trainer.logged_metrics["val_loss_epoch"].cpu().detach().item()
-        )
-        pl_module.log_store.get("val_acc_epoch").append(
-            trainer.logged_metrics["val_acc_epoch"].cpu().detach().item()
-        )
-
-
-    def on_test_epoch_end(
-        self, trainer: pl.Trainer, pl_module: pl.LightningModule
-    ) -> None:
-        super().on_test_epoch_end(trainer, pl_module)
-        print(
-            f"\nEpoch: {trainer.current_epoch}, Test Loss: {trainer.logged_metrics['test_loss_epoch']}, Test Accuracy: {trainer.logged_metrics['test_acc_epoch']}"
-        )
-        pl_module.log_store.get("test_loss_epoch").append(
-            trainer.logged_metrics["test_loss_epoch"].cpu().detach().item()
-        )
-        pl_module.log_store.get("test_acc_epoch").append(
-            trainer.logged_metrics["test_acc_epoch"].cpu().detach().item()
-        )

utilities/config.py

@@ -1,58 +0,0 @@
-# Seed
-SEED = 1
-
-# Dataset
-
-CLASSES = (
-    "Airplane",
-    "Automobile",
-    "Bird",
-    "Cat",
-    "Deer",
-    "Dog",
-    "Frog",
-    "Horse",
-    "Ship",
-    "Truck",
-)
-
-SHUFFLE = True
-DATA_DIR = "../data"
-NUM_WORKERS = 4
-PIN_MEMORY = True
-
-# Training Hyperparameters
-
-INPUT_SIZE = (3, 32, 32)
-NUM_CLASSES = 10
-LEARNING_RATE = 0.001
-WEIGHT_DECAY = 1e-4
-BATCH_SIZE = 512
-NUM_EPOCHS = 24
-DROPOUT_PERCENTAGE = 0.05
-LAYER_NORM = "bn"  # Batch Normalization
-
-# OPTIMIZER & SCHEDULER
-
-LRFINDER_END_LR = 0.1
-LRFINDER_NUM_ITERATIONS = 50
-LRFINDER_STEP_MODE = "exp"
-
-OCLR_DIV_FACTOR = 100
-OCLR_FINAL_DIV_FACTOR = 100
-OCLR_THREE_PHASE = False
-OCLR_ANNEAL_STRATEGY = "linear"
-
-# Compute Related
-
-ACCELERATOR = "cuda"
-PRECISION = 32
-
-# Store
-
-TRAINING_STAT_STORE = "Store/training_stats.csv"
-MODEL_SAVE_PATH = "Store/model.pth"
-
-# Visualization
-
-NORM_CONF_MAT = True

utilities/dataset.py

@@ -1,92 +0,0 @@
-import numpy as np
-import pytorch_lightning as pl
-import torch
-from torchvision import datasets
-
-
-class CIFAR10(torch.utils.data.Dataset):
-    def __init__(self, dataset, transform=None) -> None:
-        # Initialize dataset and transform
-        self.dataset = dataset
-        self.transform = transform
-
-    def __len__(self) -> int:
-        # Return the length of the dataset
-        return len(self.dataset)
-
-    def __getitem__(self, index):
-        # Get image and label
-        image, label = self.dataset[index]
-
-        # Convert PIL image to numpy array
-        image = np.array(image)
-
-        # Apply transformations
-        if self.transform:
-            image = self.transform(image=image)["image"]
-
-        return (image, label)
-
-
-class CIFAR10DataModule(pl.LightningDataModule):
-    def __init__(
-        self,
-        train_transforms,
-        val_transforms,
-        shuffle=True,
-        data_dir="../data",
-        batch_size=64,
-        num_workers=-1,
-        pin_memory=True,
-    ):
-        super().__init__()
-        self.shuffle = shuffle
-        self.data_dir = data_dir
-        self.batch_size = batch_size
-        self.num_workers = num_workers
-        self.pin_memory = pin_memory
-        self.train_transforms = train_transforms
-        self.val_transforms = val_transforms
-        self.train_data = None
-        self.val_data = None
-
-    def prepare_data(self):
-        datasets.CIFAR10(self.data_dir, train=True, download=True)
-        datasets.CIFAR10(self.data_dir, train=False, download=True)
-
-    def setup(self, stage):
-        self.train_data = CIFAR10(
-            datasets.CIFAR10(root=self.data_dir, train=True, download=False),
-            transform=self.train_transforms,
-        )
-        self.val_data = CIFAR10(
-            datasets.CIFAR10(root=self.data_dir, train=False, download=False),
-            transform=self.val_transforms,
-        )
-
-    def train_dataloader(self):
-        return torch.utils.data.DataLoader(
-            self.train_data,
-            batch_size=self.batch_size,
-            shuffle=self.shuffle,
-            num_workers=self.num_workers,
-            pin_memory=self.pin_memory,
-        )
-
-    def val_dataloader(self):
-        return torch.utils.data.DataLoader(
-            self.val_data,
-            batch_size=self.batch_size,
-            shuffle=False,
-            num_workers=self.num_workers,
-            pin_memory=self.pin_memory,
-        )
-
-    def test_dataloader(self):
-        return torch.utils.data.DataLoader(
-            self.val_data,
-            batch_size=self.batch_size,
-            shuffle=False,
-            num_workers=self.num_workers,
-            pin_memory=self.pin_memory,
-        )

utilities/resnet.py

@@ -1,162 +0,0 @@
-"""
-ResNet in PyTorch.
-For Pre-activation ResNet, see 'preact_resnet.py'.
-
-Reference:
-[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun
-    Deep Residual Learning for Image Recognition. arXiv:1512.03385
-"""
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import pytorch_lightning as pl
-from torchmetrics.functional import accuracy
-from torchvision import transforms
-from torch.utils.data import DataLoader
-from torchvision.datasets import CIFAR10
-import albumentations as A
-from albumentations.pytorch import ToTensorV2
-
-
-class BasicBlock(nn.Module):
-    expansion = 1
-
-    def __init__(self, in_planes, planes, stride=1):
-        super(BasicBlock, self).__init__()
-        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
-        self.bn1 = nn.BatchNorm2d(planes)
-        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
-        self.bn2 = nn.BatchNorm2d(planes)
-
-        self.shortcut = nn.Sequential()
-        if stride != 1 or in_planes != self.expansion*planes:
-            self.shortcut = nn.Sequential(
-                nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
-                nn.BatchNorm2d(self.expansion*planes)
-            )
-
-    def forward(self, x):
-        out = F.relu(self.bn1(self.conv1(x)))
-        out = self.bn2(self.conv2(out))
-        out += self.shortcut(x)
-        out = F.relu(out)
-        return out
-
-
-class LitResNet(pl.LightningModule):
-    def __init__(self, block, num_blocks, num_classes=10,batch_size=128):
-        super(LitResNet, self).__init__()
-        self.batch_size = batch_size
-        self.in_planes = 64
-        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
-        self.bn1 = nn.BatchNorm2d(64)
-        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
-        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
-        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
-        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
-        self.linear = nn.Linear(512*block.expansion, num_classes)
-
-    def _make_layer(self, block, planes, num_blocks, stride):
-        strides = [stride] + [1]*(num_blocks-1)
-        layers = []
-        for stride in strides:
-            layers.append(block(self.in_planes, planes, stride))
-            self.in_planes = planes * block.expansion
-        return nn.Sequential(*layers)
-
-        
-    def forward(self, x):
-        out = F.relu(self.bn1(self.conv1(x)))
-        out = self.layer1(out)
-        out = self.layer2(out)
-        out = self.layer3(out)
-        out = self.layer4(out)
-        out = F.avg_pool2d(out, 4)
-        out = out.view(out.size(0), -1)
-        out = self.linear(out)
-        return out
-    
-        
-    
-    def training_step(self, batch, batch_idx):
-        x, y = batch
-        y_hat = self(x)
-        # Calculate loss
-        loss = F.cross_entropy(y_hat, y)
-        #Calculate accuracy
-        acc = accuracy(y_hat, y)
-        self.log_dict(
-            {"train_loss": loss, "train_acc": acc},
-            on_step=True,
-            on_epoch=True,
-            prog_bar=True,
-            logger=True,
-        )
-        return loss
-    
-    def validation_step(self, batch, batch_idx):
-        x, y = batch
-        y_hat = self(x)
-        loss = F.cross_entropy(y_hat, y)
-        acc = accuracy(y_hat, y)
-        self.log_dict(
-            {"val_loss": loss, "val_acc": acc},
-            on_step=True,
-            on_epoch=True,
-            prog_bar=True,
-            logger=True,
-        )
-        return loss
-
-    def test_step(self, batch, batch_idx):
-        x, y = batch
-        y_hat = self(x)
-
-        argmax_pred = y_hat.argmax(dim=1).cpu()
-        loss = F.cross_entropy(y_hat, y)
-        acc = accuracy(y_hat, y)
-        self.log_dict(
-            {"test_loss": loss, "test_acc": acc},
-            on_step=True,
-            on_epoch=True,
-            prog_bar=True,
-            logger=True,
-        )
-
-        # Update the confusion matrix
-        self.confusion_matrix.update(y_hat, y)
-
-        # Store the predictions, labels and incorrect predictions
-        x, y, y_hat, argmax_pred = (
-            x.cpu(),
-            y.cpu(),
-            y_hat.cpu(),
-            argmax_pred.cpu(),
-        )
-        self.pred_store["test_preds"] = torch.cat(
-            (self.pred_store["test_preds"], argmax_pred), dim=0
-        )
-        self.pred_store["test_labels"] = torch.cat(
-            (self.pred_store["test_labels"], y), dim=0
-        )
-        for d, t, p, o in zip(x, y, argmax_pred, y_hat):
-            if p.eq(t.view_as(p)).item() == False:
-                self.pred_store["test_incorrect"].append(
-                    (d.cpu(), t, p, o[p.item()].cpu())
-                )
-
-        return loss
-
-
-def configure_optimizers(self):
-        return torch.optim.Adam(self.parameters(), lr=0.02)
-
-def LitResNet18():
-    return LitResNet(BasicBlock, [2, 2, 2, 2])
-
-def LitResNet34(): 
-    return LitResNet(BasicBlock, [3, 4, 6, 3])
-
-
-
-

utilities/transforms.py

@@ -1,20 +0,0 @@
-# Third-Party Imports
-import torch
-import albumentations as A
-from albumentations.pytorch import ToTensorV2
-
-
-# Train Phase transformations
-train_set_transforms = {
-    'randomcrop': A.RandomCrop(height=32, width=32, p=0.2),
-    'horizontalflip': A.HorizontalFlip(),
-    'cutout': A.CoarseDropout(max_holes=1, max_height=16, max_width=16, min_holes=1, min_height=1, min_width=1, fill_value=[0.49139968*255, 0.48215827*255 ,0.44653124*255], mask_fill_value=None),
-    'normalize': A.Normalize((0.49139968, 0.48215827, 0.44653124), (0.24703233, 0.24348505, 0.26158768)),
-    'standardize': ToTensorV2(),
-}
-
-# Test Phase transformations
-test_set_transforms = {
-    'normalize': A.Normalize((0.49139968, 0.48215827, 0.44653124), (0.24703233, 0.24348505, 0.26158768)),
-    'standardize': ToTensorV2()
-}

utilities/visualise.py

@@ -1,78 +0,0 @@
-import matplotlib.pyplot as plt
-from torchvision import transforms
-
-
-def plot_class_label_counts(data_loader, classes):
-    class_counts = {}
-    for class_name in classes:
-        class_counts[class_name] = 0
-    for _, batch_label in data_loader:
-        for label in batch_label:
-            class_counts[classes[label.item()]] += 1
-
-    fig = plt.figure()
-    plt.suptitle("Class Distribution")
-    plt.bar(range(len(class_counts)), list(class_counts.values()))
-    plt.xticks(range(len(class_counts)), list(class_counts.keys()), rotation=90)
-    plt.tight_layout()
-    plt.show()
-
-
-def plot_data_samples(data_loader, classes):
-    batch_data, batch_label = next(iter(data_loader))
-
-    fig = plt.figure()
-    plt.suptitle("Data Samples with Labels post Transforms")
-    for i in range(12):
-        plt.subplot(3, 4, i + 1)
-        plt.tight_layout()
-        # unnormalize = T.Normalize((-mean / std).tolist(), (1.0 / std).tolist())
-        unnormalized = transforms.Normalize(
-            (-1.98947368, -1.98436214, -1.71072797), (4.048583, 4.11522634, 3.83141762)
-        )(batch_data[i])
-        plt.imshow(transforms.ToPILImage()(unnormalized))
-        plt.title(
-            classes[batch_label[i].item()],
-        )
-
-        plt.xticks([])
-        plt.yticks([])
-
-
-def plot_model_training_curves(train_accs, test_accs, train_losses, test_losses):
-    fig, axs = plt.subplots(2, 2, figsize=(15, 10))
-    axs[0, 0].plot(train_losses)
-    axs[0, 0].set_title("Training Loss")
-    axs[1, 0].plot(train_accs)
-    axs[1, 0].set_title("Training Accuracy")
-    axs[0, 1].plot(test_losses)
-    axs[0, 1].set_title("Test Loss")
-    axs[1, 1].plot(test_accs)
-    axs[1, 1].set_title("Test Accuracy")
-    plt.plot()
-
-
-def plot_incorrect_preds(incorrect, classes, num_imgs):
-    # num_imgs is a multiple of 5
-    assert num_imgs % 5 == 0
-    assert len(incorrect) >= num_imgs
-
-    # incorrect (data, target, pred, output)
-    print(f"Total Incorrect Predictions {len(incorrect)}")
-    fig = plt.figure(figsize=(10, num_imgs // 2))
-    plt.suptitle("Target | Predicted Label")
-    for i in range(num_imgs):
-        plt.subplot(num_imgs // 5, 5, i + 1, aspect="auto")
-
-        # unnormalize = T.Normalize((-mean / std).tolist(), (1.0 / std).tolist())
-        unnormalized = transforms.Normalize(
-            (-1.98947368, -1.98436214, -1.71072797), (4.048583, 4.11522634, 3.83141762)
-        )(incorrect[i][0])
-        plt.imshow(transforms.ToPILImage()(unnormalized))
-        plt.title(
-            f"{classes[incorrect[i][1].item()]}|{classes[incorrect[i][2].item()]}",
-            # fontsize=8,
-        )
-        plt.xticks([])
-        plt.yticks([])
-    plt.tight_layout()