Upload 8 files

augmentations.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python3
+"""
+Function used for visualization of data and results
+Author: Shilpaj Bhalerao
+Date: Jul 23, 2023
+"""
+# 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()
+}
+
+
+class AddGaussianNoise(object):
+    """
+    Class for custom augmentation strategy
+    """
+    def __init__(self, mean=0., std=1.):
+        """
+        Constructor
+        """
+        self.std = std
+        self.mean = mean
+
+    def __call__(self, tensor):
+        """
+        Augmentation strategy to be implemented when called
+        """
+        return tensor + torch.randn(tensor.size()) * self.std + self.mean
+
+    def __repr__(self):
+        """
+        Method to print more infor about the strategy
+        """
+        return f"{self.__class__.__name__}(mean={self.mean}, std={self.std})"
+
+# Usage details
+# transforms = transforms.Compose([
+#     transforms.ToTensor(),
+#     AddGaussianNoise(0., 1.0),
+#     ])

datasets.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+"""
+Module containing wrapper classes for PyTorch Datasets
+Author: Shilpaj Bhalerao
+Date: Jun 25, 2023
+"""
+# Standard Library Imports
+from typing import Tuple
+
+# Third-Party Imports
+from torchvision import datasets, transforms
+
+
+class AlbumDataset(datasets.CIFAR10):
+    """
+    Wrapper class to use albumentations library with PyTorch Dataset
+    """
+    def __init__(self, root: str = "./data", train: bool = True, download: bool = True, transform: list = None):
+        """
+        Constructor
+        :param root: Directory at which data is stored
+        :param train: Param to distinguish if data is training or test
+        :param download: Param to download the dataset from source
+        :param transform: List of transformation to be performed on the dataset
+        """
+        super().__init__(root=root, train=train, download=download, transform=transform)
+
+    def __getitem__(self, index: int) -> Tuple:
+        """
+        Method to return image and its label
+        :param index: Index of image and label in the dataset
+        """
+        image, label = self.data[index], self.targets[index]
+
+        if self.transform:
+            transformed = self.transform(image=image)
+            image = transformed["image"]
+        return image, label

gradio_utils.py

@@ -0,0 +1,151 @@
+# Import all the required modules
+import os
+os.environ['KMP_DUPLICATE_LIB_OK']='True'
+import math
+from collections import OrderedDict
+import sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torchvision import datasets, transforms
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+from torch_lr_finder import LRFinder
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
+from pytorch_grad_cam.utils.image import show_cam_on_image
+from utils import display_gradcam_output
+from torchmetrics import Accuracy
+from torchvision.datasets import CIFAR10 
+import torch
+from pytorch_lightning import LightningModule, Trainer
+from torch import nn
+from torch.nn import functional as F
+from torch.utils.data import DataLoader, random_split
+from torchmetrics import Accuracy
+from torchvision import transforms
+import matplotlib.pyplot as plt
+import random
+from resnet import *
+
+PATH_DATASETS = os.environ.get("PATH_DATASETS", ".")
+classes = ('plane', 'car', 'bird', 'cat', 'deer',
+           'dog', 'frog', 'horse', 'ship', 'truck')
+
+
+
+class LitCIFAR(LightningModule):
+    def __init__(self, data_dir=PATH_DATASETS, hidden_size=16, learning_rate=2e-4):
+
+        super().__init__()
+
+        # Set our init args as class attributes
+        self.data_dir = data_dir
+        self.hidden_size = hidden_size
+        self.learning_rate = learning_rate
+
+        # Hardcode some dataset specific attributes
+        self.num_classes = 10
+
+        self.misclassified_indices = []
+
+        # Define PyTorch model
+        self.model = ResNet18()
+
+        self.accuracy = Accuracy(num_classes=self.num_classes, task='multiclass')
+
+    def forward(self, x):
+        x = self.model(x)
+        x = x.view(-1, 10)
+        return F.log_softmax(x, dim=1)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        logits = self(x)
+        loss = F.nll_loss(logits, y)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        logits = self(x)
+        loss = F.nll_loss(logits, y)
+        preds = torch.argmax(logits, dim=1)
+        self.accuracy(preds, y)
+        self.log("val_loss", loss, prog_bar=True)
+        self.log("val_acc", self.accuracy, prog_bar=True)
+        return loss
+
+    # def test_step(self, batch, batch_idx):
+    #     # Here we just reuse the validation_step for testing
+    #     return self.validation_step(batch, batch_idx)
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
+        return optimizer
+
+
+modelfin = LitCIFAR()
+lit_cifar_instance = LitCIFAR()
+# Load the state dictionary from the checkpoint file
+modelfin.load_state_dict(torch.load("submission_gradio/models/model.ckpt"))
+# Set the model to evaluation mode
+modelfin.eval()
+# If you need to use the model on a GPU, move the model to GPU after loading the state dict and setting it to eval mode
+modelfin = modelfin.to('cuda')
+
+inv_normalize = transforms.Normalize(
+    mean=[-1.9899, -1.9844, -1.7111],
+    std=[4.0486, 4.1152, 3.8314]
+)
+
+from torch.utils.data import Dataset
+import numpy as np
+class CIFAR10Dataset(Dataset):
+    def __init__(self, data_dir, train=True, transform=None):
+        self.data = CIFAR10(data_dir, train=train, download=True, transform=None)
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        image, label = self.data[idx]
+        if self.transform:
+            # Apply the transformation in the correct format
+            image = self.transform(image=np.array(image))['image']
+        return image, label
+
+transform = A.Compose(
+    [
+        A.RandomCrop(height=32, width=32, p=0.2),
+        A.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261)),
+        A.HorizontalFlip(),
+        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),
+        ToTensorV2(),
+    ]
+)
+test_transform = A.Compose(
+    [
+        A.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261)),
+        ToTensorV2(),
+    ]
+)
+
+cifar_full = CIFAR10Dataset(PATH_DATASETS, train=True, transform=transform)
+cifar_test = CIFAR10Dataset(PATH_DATASETS, train=False, transform=test_transform)
+classes = ('plane', 'car', 'bird', 'cat', 'deer',
+           'dog', 'frog', 'horse', 'ship', 'truck')
+
+test_loader = DataLoader(cifar_test, batch_size=512, num_workers=os.cpu_count())
+
+import gradio as gr
+from utils import display_gradcam_output
+from utils import get_misclassified_data
+from visualize import display_cifar_misclassified_data
+from PIL import Image
+
+#misclassified_data, classes, inv_normalize, modelfin, target_layers, targets, number_of_samples=2, transparency=0.7
+
+
+

resnet.py

@@ -0,0 +1,76 @@
+"""
+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.nn as nn
+import torch.nn.functional as F
+
+
+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 ResNet(nn.Module):
+    def __init__(self, block, num_blocks, num_classes=10):
+        super(ResNet, self).__init__()
+        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 ResNet18():
+    return ResNet(BasicBlock, [2, 2, 2, 2])
+
+
+def ResNet34():
+    return ResNet(BasicBlock, [3, 4, 6, 3])

training_utils.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python3
+"""
+Utilities for Model Training
+Author: Shilpaj Bhalerao
+Date: Jun 21, 2023
+"""
+# Standard Library Imports
+
+# Third-Party Imports
+from tqdm import tqdm
+import torch
+
+
+def get_correct_predictions(prediction, labels):
+    """
+    Function to return total number of correct predictions
+    :param prediction: Model predictions on a given sample of data
+    :param labels: Correct labels of a given sample of data
+    :return: Number of correct predictions
+    """
+    return prediction.argmax(dim=1).eq(labels).sum().item()
+
+
+def train(model, device, train_loader, optimizer, criterion, scheduler=None):
+    """
+    Function to train model on the training dataset
+    :param model: Model architecture
+    :param device: Device on which training is to be done (GPU/CPU)
+    :param train_loader: DataLoader for training dataset
+    :param optimizer: Optimization algorithm to be used for updating weights
+    :param criterion: Loss function for training
+    :param scheduler: Scheduler for learning rate
+    """
+    # Enable layers like Dropout for model training
+    model.train()
+
+    # Utility to display training progress
+    pbar = tqdm(train_loader)
+
+    # Variables to track loss and accuracy during training
+    train_loss = 0
+    correct = 0
+    processed = 0
+
+    # Iterate over each batch and fetch images and labels from the batch
+    for batch_idx, (data, target) in enumerate(pbar):
+
+        # Put the images and labels on the selected device
+        data, target = data.to(device), target.to(device)
+
+        # Reset the gradients for each batch
+        optimizer.zero_grad()
+
+        # Predict
+        pred = model(data)
+
+        # Calculate loss
+        loss = criterion(pred, target)
+        train_loss += loss.item()
+
+        # Backpropagation
+        loss.backward()
+        optimizer.step()
+
+        # Use learning rate scheduler if defined
+        if scheduler:
+            scheduler.step()
+
+        # Get total number of correct predictions
+        correct += get_correct_predictions(pred, target)
+        processed += len(data)
+
+        # Display the training information
+        pbar.set_description(
+            desc=f'Train: Loss={loss.item():0.4f} Batch_id={batch_idx} Accuracy={100 * correct / processed:0.2f}')
+
+    return correct, processed, train_loss
+
+
+def test(model, device, test_loader, criterion):
+    """
+    Function to test the model training progress on the test dataset
+    :param model: Model architecture
+    :param device: Device on which training is to be done (GPU/CPU)
+    :param test_loader: DataLoader for test dataset
+    :param criterion: Loss function for test dataset
+    """
+    # Disable layers like Dropout for model inference
+    model.eval()
+
+    # Variables to track loss and accuracy
+    test_loss = 0
+    correct = 0
+
+    # Disable gradient updation
+    with torch.no_grad():
+        # Iterate over each batch and fetch images and labels from the batch
+        for batch_idx, (data, target) in enumerate(test_loader):
+
+            # Put the images and labels on the selected device
+            data, target = data.to(device), target.to(device)
+
+            # Pass the images to the output and get the model predictions
+            output = model(data)
+            test_loss += criterion(output, target).item()  # sum up batch loss
+
+            # Sum up batch correct predictions
+            correct += get_correct_predictions(output, target)
+
+    # Calculate test loss for a epoch
+    test_loss /= len(test_loader.dataset)
+
+    print('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format(
+        test_loss, correct, len(test_loader.dataset),
+        100. * correct / len(test_loader.dataset)))
+
+    return correct, test_loss
+
+
+def get_lr(optimizer):
+    """
+    Function to track learning rate while model training
+    :param optimizer: Optimizer used for training
+    """
+    for param_group in optimizer.param_groups:
+        return param_group['lr']

utils.py

@@ -0,0 +1,216 @@
+#!/usr/bin/env python3
+"""
+Utility Script containing functions to be used for training
+Author: Shilpaj Bhalerao
+"""
+# Standard Library Imports
+import math
+from typing import NoReturn
+import io
+from PIL import Image
+# Third-Party Imports
+import numpy as np
+import matplotlib.pyplot as plt
+import torch
+from torchsummary import summary
+from torchvision import transforms
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+
+
+def get_summary(model, input_size: tuple) -> NoReturn:
+    """
+    Function to get the summary of the model architecture
+    :param model: Object of model architecture class
+    :param input_size: Input data shape (Channels, Height, Width)
+    """
+    use_cuda = torch.cuda.is_available()
+    device = torch.device("cuda" if use_cuda else "cpu")
+    network = model.to(device)
+    summary(network, input_size=input_size)
+
+
+def get_misclassified_data(model, device, test_loader):
+    """
+    Function to run the model on test set and return misclassified images
+    :param model: Network Architecture
+    :param device: CPU/GPU
+    :param test_loader: DataLoader for test set
+    """
+    # Prepare the model for evaluation i.e. drop the dropout layer
+    model.eval()
+
+    # List to store misclassified Images
+    misclassified_data = []
+
+    # Reset the gradients
+    with torch.no_grad():
+        # Extract images, labels in a batch
+        for data, target in test_loader:
+
+            # Migrate the data to the device
+            data, target = data.to(device), target.to(device)
+
+            # Extract single image, label from the batch
+            for image, label in zip(data, target):
+
+                # Add batch dimension to the image
+                image = image.unsqueeze(0)
+
+                # Get the model prediction on the image
+                output = model(image)
+
+                # Convert the output from one-hot encoding to a value
+                pred = output.argmax(dim=1, keepdim=True)
+
+                # If prediction is incorrect, append the data
+                if pred != label:
+                    misclassified_data.append((image, label, pred))
+    return misclassified_data
+
+
+# -------------------- DATA STATISTICS --------------------
+def get_mnist_statistics(data_set, data_set_type='Train'):
+    """
+    Function to return the statistics of the training data
+    :param data_set: Training dataset
+    :param data_set_type: Type of dataset [Train/Test/Val]
+    """
+    # We'd need to convert it into Numpy! Remember above we have converted it into tensors already
+    train_data = data_set.train_data
+    train_data = data_set.transform(train_data.numpy())
+
+    print(f'[{data_set_type}]')
+    print(' - Numpy Shape:', data_set.train_data.cpu().numpy().shape)
+    print(' - Tensor Shape:', data_set.train_data.size())
+    print(' - min:', torch.min(train_data))
+    print(' - max:', torch.max(train_data))
+    print(' - mean:', torch.mean(train_data))
+    print(' - std:', torch.std(train_data))
+    print(' - var:', torch.var(train_data))
+
+    dataiter = next(iter(data_set))
+    images, labels = dataiter[0], dataiter[1]
+
+    print(images.shape)
+    print(labels)
+
+    # Let's visualize some of the images
+    plt.imshow(images[0].numpy().squeeze(), cmap='gray')
+
+
+def get_cifar_property(images, operation):
+    """
+    Get the property on each channel of the CIFAR
+    :param images: Get the property value on the images
+    :param operation: Mean, std, Variance, etc
+    """
+    param_r = eval('images[:, 0, :, :].' + operation + '()')
+    param_g = eval('images[:, 1, :, :].' + operation + '()')
+    param_b = eval('images[:, 2, :, :].' + operation + '()')
+    return param_r, param_g, param_b
+
+
+def get_cifar_statistics(data_set, data_set_type='Train'):
+    """
+    Function to get the statistical information of the CIFAR dataset
+    :param data_set: Training set of CIFAR
+    :param data_set_type: Training or Test data
+    """
+    # Images in the dataset
+    images = [item[0] for item in data_set]
+    images = torch.stack(images, dim=0).numpy()
+
+    # Calculate mean over each channel
+    mean_r, mean_g, mean_b = get_cifar_property(images, 'mean')
+
+    # Calculate Standard deviation over each channel
+    std_r, std_g, std_b = get_cifar_property(images, 'std')
+
+    # Calculate min value over each channel
+    min_r, min_g, min_b = get_cifar_property(images, 'min')
+
+    # Calculate max value over each channel
+    max_r, max_g, max_b = get_cifar_property(images, 'max')
+
+    # Calculate variance value over each channel
+    var_r, var_g, var_b = get_cifar_property(images, 'var')
+
+    print(f'[{data_set_type}]')
+    print(f' - Total {data_set_type} Images: {len(data_set)}')
+    print(f' - Tensor Shape: {images[0].shape}')
+    print(f' - min: {min_r, min_g, min_b}')
+    print(f' - max: {max_r, max_g, max_b}')
+    print(f' - mean: {mean_r, mean_g, mean_b}')
+    print(f' - std: {std_r, std_g, std_b}')
+    print(f' - var: {var_r, var_g, var_b}')
+
+    # Let's visualize some of the images
+    plt.imshow(np.transpose(images[1].squeeze(), (1, 2, 0)))
+
+
+# -------------------- GradCam --------------------
+def display_gradcam_output(data: list,
+                           classes,
+                           inv_normalize: transforms.Normalize,
+                           model,
+                           target_layers,
+                           targets=None,
+                           number_of_samples: int = 10,
+                           transparency: float = 0.60):
+    """
+    Function to visualize GradCam output on the data
+    :param data: List[Tuple(image, label)]
+    :param classes: Name of classes in the dataset
+    :param inv_normalize: Mean and Standard deviation values of the dataset
+    :param model: Model architecture
+    :param target_layers: Layers on which GradCam should be executed
+    :param targets: Classes to be focused on for GradCam
+    :param number_of_samples: Number of images to print
+    :param transparency: Weight of Normal image when mixed with activations
+    """
+    # Plot configuration
+    fig = plt.figure(figsize=(10, 10))
+    x_count = 5
+    y_count = math.ceil(number_of_samples / x_count) 
+
+    # Create an object for GradCam
+    cam = GradCAM(model=model, target_layers=target_layers)
+
+    # Iterate over number of specified images
+    for i in range(number_of_samples):
+        plt.subplot(y_count, x_count, i + 1)
+        input_tensor = data[i][0]
+
+        # Get the activations of the layer for the images
+        grayscale_cam = cam(input_tensor=input_tensor, targets=targets)
+        grayscale_cam = grayscale_cam[0, :]
+
+        # Get back the original image
+        img = input_tensor.squeeze(0).to('cpu')
+        img = inv_normalize(img)
+        rgb_img = np.transpose(img, (1, 2, 0))
+        rgb_img = rgb_img.numpy().astype(np.float32)
+
+# Ensure the image data is within the [0, 1] range
+        rgb_img = np.clip(rgb_img, 0, 1)
+
+        # Mix the activations on the original image
+        visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True, image_weight=transparency)
+
+        # Display the images on the plot
+        plt.imshow(visualization)
+        plt.title(r"Correct: " + classes[data[i][1].item()] + '\n' + 'Output: ' + classes[data[i][2].item()])
+        plt.xticks([])
+        plt.yticks([])
+
+    plt.tight_layout()
+
+    # Save the entire figure to a BytesIO object
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    img_var = Image.open(buf)
+
+    return img_var
+

visualize.py

@@ -0,0 +1,396 @@
+#!/usr/bin/env python3
+"""
+Function used for visualization of data and results
+Author: Shilpaj Bhalerao
+Date: Jun 21, 2023
+"""
+# Standard Library Imports
+import math
+from dataclasses import dataclass
+from typing import NoReturn
+import io
+from PIL import Image
+# Third-Party Imports
+import numpy as np
+import matplotlib.pyplot as plt
+import pandas as pd
+import seaborn as sn
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from sklearn.metrics import confusion_matrix
+
+
+# ---------------------------- DATA SAMPLES ----------------------------
+def display_mnist_data_samples(dataset, number_of_samples: int) -> NoReturn:
+    """
+    Function to display samples for dataloader
+    :param dataset: Train or Test dataset transformed to Tensor
+    :param number_of_samples: Number of samples to be displayed
+    """
+    # Get batch from the data_set
+    batch_data = []
+    batch_label = []
+    for count, item in enumerate(dataset):
+        if not count <= number_of_samples:
+            break
+        batch_data.append(item[0])
+        batch_label.append(item[1])
+
+    # Plot the samples from the batch
+    fig = plt.figure()
+    x_count = 5
+    y_count = 1 if number_of_samples <= 5 else math.floor(number_of_samples / x_count)
+
+    # Plot the samples from the batch
+    for i in range(number_of_samples):
+        plt.subplot(y_count, x_count, i + 1)
+        plt.tight_layout()
+        plt.imshow(batch_data[i].squeeze(), cmap='gray')
+        plt.title(batch_label[i])
+        plt.xticks([])
+        plt.yticks([])
+
+
+def display_cifar_data_samples(data_set, number_of_samples: int, classes: list):
+    """
+    Function to display samples for data_set
+    :param data_set: Train or Test data_set transformed to Tensor
+    :param number_of_samples: Number of samples to be displayed
+    :param classes: Name of classes to be displayed
+    """
+    # Get batch from the data_set
+    batch_data = []
+    batch_label = []
+    for count, item in enumerate(data_set):
+        if not count <= number_of_samples:
+            break
+        batch_data.append(item[0])
+        batch_label.append(item[1])
+    batch_data = torch.stack(batch_data, dim=0).numpy()
+
+    # Plot the samples from the batch
+    fig = plt.figure()
+    x_count = 5
+    y_count = 1 if number_of_samples <= 5 else math.floor(number_of_samples / x_count)
+
+    for i in range(number_of_samples):
+        plt.subplot(y_count, x_count, i + 1)
+        plt.tight_layout()
+        plt.imshow(np.transpose(batch_data[i].squeeze(), (1, 2, 0)))
+        plt.title(classes[batch_label[i]])
+        plt.xticks([])
+        plt.yticks([])
+
+
+# ---------------------------- MISCLASSIFIED DATA ----------------------------
+def display_cifar_misclassified_data(data: list,
+                                     classes,
+                                     inv_normalize: transforms.Normalize,
+                                     number_of_samples: int = 10):
+    """
+    Function to plot images with labels
+    :param data: List[Tuple(image, label)]
+    :param classes: Name of classes in the dataset
+    :param inv_normalize: Mean and Standard deviation values of the dataset
+    :param number_of_samples: Number of images to print
+    """
+    fig = plt.figure(figsize=(10, 10))
+
+    x_count = 5
+    y_count = math.ceil(number_of_samples / x_count)
+
+    for i in range(number_of_samples):
+        plt.subplot(y_count, x_count, i + 1)
+        img = data[i][0].squeeze().to('cpu')
+        img = inv_normalize(img)
+        plt.imshow(np.transpose(img, (1, 2, 0)))
+        plt.title(r"Correct: " + classes[data[i][1].item()] + '\n' + 'Output: ' + classes[data[i][2].item()])
+        plt.xticks([])
+        plt.yticks([])
+
+    plt.tight_layout()
+
+    # Save the entire figure to a BytesIO object
+    buftwo = io.BytesIO()
+    plt.savefig(buftwo, format='png')
+    buftwo.seek(0)
+    img_var = Image.open(buftwo)
+
+    # Optional: Close the buffer
+
+    return img_var
+
+def display_mnist_misclassified_data(data: list,
+                                     number_of_samples: int = 10):
+    """
+    Function to plot images with labels
+    :param data: List[Tuple(image, label)]
+    :param number_of_samples: Number of images to print
+    """
+    fig = plt.figure(figsize=(8, 5))
+
+    x_count = 5
+    y_count = 1 if number_of_samples <= 5 else math.floor(number_of_samples / x_count)
+
+    for i in range(number_of_samples):
+        plt.subplot(y_count, x_count, i + 1)
+        img = data[i][0].squeeze(0).to('cpu')
+        plt.imshow(np.transpose(img, (1, 2, 0)), cmap='gray')
+        plt.title(r"Correct: " + str(data[i][1].item()) + '\n' + 'Output: ' + str(data[i][2].item()))
+        plt.xticks([])
+        plt.yticks([])
+
+
+# ---------------------------- AUGMENTATION SAMPLES ----------------------------
+def visualize_cifar_augmentation(data_set, data_transforms):
+    """
+    Function to visualize the augmented data
+    :param data_set: Dataset without transformations
+    :param data_transforms: Dictionary of transforms
+    """
+    sample, label = data_set[6]
+    total_augmentations = len(data_transforms)
+
+    fig = plt.figure(figsize=(10, 5))
+    for count, (key, trans) in enumerate(data_transforms.items()):
+        if count == total_augmentations - 1:
+            break
+        plt.subplot(math.ceil(total_augmentations / 5), 5, count + 1)
+        augmented = trans(image=sample)['image']
+        plt.imshow(augmented)
+        plt.title(key)
+        plt.xticks([])
+        plt.yticks([])
+
+
+def visualize_mnist_augmentation(data_set, data_transforms):
+    """
+    Function to visualize the augmented data
+    :param data_set: Dataset to visualize the augmentations
+    :param data_transforms: Dictionary of transforms
+    """
+    sample, label = data_set[6]
+    total_augmentations = len(data_transforms)
+
+    fig = plt.figure(figsize=(10, 5))
+    for count, (key, trans) in enumerate(data_transforms.items()):
+        if count == total_augmentations - 1:
+            break
+        plt.subplot(math.ceil(total_augmentations / 5), 5, count + 1)
+        img = trans(sample).to('cpu')
+        plt.imshow(np.transpose(img, (1, 2, 0)), cmap='gray')
+        plt.title(key)
+        plt.xticks([])
+        plt.yticks([])
+
+
+# ---------------------------- LOSS AND ACCURACIES ----------------------------
+def display_loss_and_accuracies(train_losses: list,
+                                train_acc: list,
+                                test_losses: list,
+                                test_acc: list,
+                                plot_size: tuple = (10, 10)) -> NoReturn:
+    """
+    Function to display training and test information(losses and accuracies)
+    :param train_losses: List containing training loss of each epoch
+    :param train_acc: List containing training accuracy of each epoch
+    :param test_losses: List containing test loss of each epoch
+    :param test_acc: List containing test accuracy of each epoch
+    :param plot_size: Size of the plot
+    """
+    # Create a plot of 2x2 of size
+    fig, axs = plt.subplots(2, 2, figsize=plot_size)
+
+    # Plot the training loss and accuracy for each epoch
+    axs[0, 0].plot(train_losses)
+    axs[0, 0].set_title("Training Loss")
+    axs[1, 0].plot(train_acc)
+    axs[1, 0].set_title("Training Accuracy")
+
+    # Plot the test loss and accuracy for each epoch
+    axs[0, 1].plot(test_losses)
+    axs[0, 1].set_title("Test Loss")
+    axs[1, 1].plot(test_acc)
+    axs[1, 1].set_title("Test Accuracy")
+
+
+# ---------------------------- Feature Maps and Kernels ----------------------------
+
+@dataclass
+class ConvLayerInfo:
+    """
+    Data Class to store Conv layer's information
+    """
+    layer_number: int
+    weights: torch.nn.parameter.Parameter
+    layer_info: torch.nn.modules.conv.Conv2d
+
+
+class FeatureMapVisualizer:
+    """
+    Class to visualize Feature Map of the Layers
+    """
+
+    def __init__(self, model):
+        """
+        Contructor
+        :param model: Model Architecture
+        """
+        self.conv_layers = []
+        self.outputs = []
+        self.layerwise_kernels = None
+
+        # Disect the model
+        counter = 0
+        model_children = model.children()
+        for children in model_children:
+            if type(children) == nn.Sequential:
+                for child in children:
+                    if type(child) == nn.Conv2d:
+                        counter += 1
+                        self.conv_layers.append(ConvLayerInfo(layer_number=counter,
+                                                              weights=child.weight,
+                                                              layer_info=child)
+                                                )
+
+    def get_model_weights(self):
+        """
+        Method to get the model weights
+        """
+        model_weights = [layer.weights for layer in self.conv_layers]
+        return model_weights
+
+    def get_conv_layers(self):
+        """
+        Get the convolution layers
+        """
+        conv_layers = [layer.layer_info for layer in self.conv_layers]
+        return conv_layers
+
+    def get_total_conv_layers(self) -> int:
+        """
+        Get total number of convolution layers
+        """
+        out = self.get_conv_layers()
+        return len(out)
+
+    def feature_maps_of_all_kernels(self, image: torch.Tensor) -> dict:
+        """
+        Get feature maps from all the kernels of all the layers
+        :param image: Image to be passed to the network
+        """
+        image = image.unsqueeze(0)
+        image = image.to('cpu')
+
+        outputs = {}
+
+        layers = self.get_conv_layers()
+        for index, layer in enumerate(layers):
+            image = layer(image)
+            outputs[str(layer)] = image
+        self.outputs = outputs
+        return outputs
+
+    def visualize_feature_map_of_kernel(self, image: torch.Tensor, kernel_number: int) -> None:
+        """
+        Function to visualize feature map of kernel number from each layer
+        :param image: Image passed to the network
+        :param kernel_number: Number of kernel in each layer (Should be less than or equal to the minimum number of kernel in the network)
+        """
+        # List to store processed feature maps
+        processed = []
+
+        # Get feature maps from all kernels of all the conv layers
+        outputs = self.feature_maps_of_all_kernels(image)
+
+        # Extract the n_th kernel's output from each layer and convert it to grayscale
+        for feature_map in outputs.values():
+            try:
+                feature_map = feature_map[0][kernel_number]
+            except IndexError:
+                print("Filter number should be less than the minimum number of channels in a network")
+                break
+            finally:
+                gray_scale = feature_map / feature_map.shape[0]
+                processed.append(gray_scale.data.numpy())
+
+        # Plot the Feature maps with layer and kernel number
+        x_range = len(outputs) // 5 + 4
+        fig = plt.figure(figsize=(10, 10))
+        for i in range(len(processed)):
+            a = fig.add_subplot(x_range, 5, i + 1)
+            imgplot = plt.imshow(processed[i])
+            a.axis("off")
+            title = f"{list(outputs.keys())[i].split('(')[0]}_l{i + 1}_k{kernel_number}"
+            a.set_title(title, fontsize=10)
+
+    def get_max_kernel_number(self):
+        """
+        Function to get maximum number of kernels in the network (for a layer)
+        """
+        layers = self.get_conv_layers()
+        channels = [layer.out_channels for layer in layers]
+        self.layerwise_kernels = channels
+        return max(channels)
+
+    def visualize_kernels_from_layer(self, layer_number: int):
+        """
+        Visualize Kernels from a layer
+        :param layer_number: Number of layer from which kernels are to be visualized
+        """
+        # Get the kernels number for each layer
+        self.get_max_kernel_number()
+
+        # Zero Indexing
+        layer_number = layer_number - 1
+        _kernels = self.layerwise_kernels[layer_number]
+
+        grid = math.ceil(math.sqrt(_kernels))
+
+        plt.figure(figsize=(5, 4))
+        model_weights = self.get_model_weights()
+        _layer_weights = model_weights[layer_number].cpu()
+        for i, filter in enumerate(_layer_weights):
+            plt.subplot(grid, grid, i + 1)
+            plt.imshow(filter[0, :, :].detach(), cmap='gray')
+            plt.axis('off')
+        plt.show()
+
+
+# ---------------------------- Confusion Matrix ----------------------------
+def visualize_confusion_matrix(classes, device: str, model,
+                               test_loader: torch.utils.data.DataLoader):
+    """
+    Function to generate and visualize confusion matrix
+    :param classes: List of class names
+    :param device: cuda/cpu
+    :param model: Model Architecture
+    :param test_loader: DataLoader for test set
+    """
+    nb_classes = len(classes)
+    device = 'cuda'
+    cm = torch.zeros(nb_classes, nb_classes)
+
+    model.eval()
+    with torch.no_grad():
+        for inputs, labels in test_loader:
+            inputs = inputs.to(device)
+            labels = labels.to(device)
+            model = model.to(device)
+
+            preds = model(inputs)
+            preds = preds.argmax(dim=1)
+
+        for t, p in zip(labels.view(-1), preds.view(-1)):
+            cm[t, p] = cm[t, p] + 1
+
+    # Build confusion matrix
+    labels = labels.to('cpu')
+    preds = preds.to('cpu')
+    cf_matrix = confusion_matrix(labels, preds)
+    df_cm = pd.DataFrame(cf_matrix / np.sum(cf_matrix, axis=1)[:, None],
+                         index=[i for i in classes],
+                         columns=[i for i in classes])
+    plt.figure(figsize=(12, 7))
+    sn.heatmap(df_cm, annot=True)