Upload 4 files

epoch=19-step=3920.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b3c28bc1f895ab927922a707c81ea52aac3bf6c311ba91577f706e997a66f73
+size 89490895

resnet_lightning.py

@@ -0,0 +1,179 @@
+# 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
+
+
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+
+from torch_lr_finder import LRFinder
+
+from pytorch_grad_cam import GradCAM
+from utils import *
+
+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 ResNet18Model(LightningModule):
+    def __init__(self, data_dir=PATH_DATASETS, block=BasicBlock, num_blocks=[2, 2, 2, 2], num_classes=10):
+        super(ResNet18Model, self).__init__()
+        self.data_dir = data_dir
+        self.num_classes = num_classes
+
+        means = [0.4914, 0.4822, 0.4465]
+        stds = [0.2470, 0.2435, 0.2616]
+
+        self.train_transforms = A.Compose(
+            [
+                A.Normalize(mean=means, std=stds, always_apply=True),
+                A.PadIfNeeded(min_height=36, min_width=36, always_apply=True),
+                A.RandomCrop(height=32, width=32, always_apply=True),
+                A.HorizontalFlip(),
+                A.CoarseDropout(max_holes=1, max_height=16, max_width=16, min_holes=1, min_height=8, min_width=8, fill_value=means),
+                ToTensorV2(),
+            ]
+        )
+        self.test_transforms = A.Compose(
+            [
+                A.Normalize(mean=means, std=stds, always_apply=True),
+                ToTensorV2(),
+            ]
+        )
+
+        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)
+
+        self.accuracy = Accuracy(task="multiclass", num_classes=10)
+
+    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
+        loss = F.cross_entropy(self(x), 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)
+
+      # Calling self.log will surface up scalars for you in TensorBoard
+      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):
+      LEARNING_RATE = 0.03
+      WEIGHT_DECAY = 1e-4
+      # # Loss Function
+      # criterion = nn.CrossEntropyLoss()
+
+      # optimizer = optim.SGD(self.parameters(), lr=LEARNING_RATE, momentum=0.9, weight_decay=WEIGHT_DECAY)
+
+
+      # lr_finder2 = LRFinder(self, optimizer, criterion, device='cuda')
+      # lr_finder2.range_test(train_loader, end_lr=10, num_iter=200, step_mode="exp")
+      # lr_finder2.plot()
+      # suggested_lr = lr_finder2.suggest_lr()
+      # lr_finder2.reset()
+      # EPOCHS = 20
+      # STEPS_PER_EPOCH = 2000
+
+      # scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer,
+      #                                           max_lr=suggested_lr,
+      #                                           steps_per_epoch=STEPS_PER_EPOCH,
+      #                                           epochs=EPOCHS,
+      #                                           pct_start=int(0.3*EPOCHS)/EPOCHS if EPOCHS != 1 else 0.5,   # 30% of total number of Epochs
+      #                                           div_factor=100,
+      #                                           three_phase=False,
+      #                                           final_div_factor=100,
+      #                                           anneal_strategy="linear"
+      #                                           )
+      return torch.optim.SGD(self.parameters(), lr=LEARNING_RATE, momentum=0.9, weight_decay=WEIGHT_DECAY)
+      # return scheduler
+
+    def prepare_data(self):
+        # download
+        Cifar10SearchDataset(self.data_dir, train=True, download=True)
+        Cifar10SearchDataset(self.data_dir, train=False, download=True)
+
+    def setup(self, stage=None):
+
+      # Assign train/val datasets for use in dataloaders
+      if stage == "fit" or stage is None:
+          cifar_full = Cifar10SearchDataset(self.data_dir, train=True, transform=self.train_transforms)
+          self.cifar_train, self.cifar_val = random_split(cifar_full, [45000, 5000])
+
+      # Assign test dataset for use in dataloader(s)
+      if stage == "test" or stage is None:
+          self.cifar_test = Cifar10SearchDataset(self.data_dir, train=False, transform=self.test_transforms)
+
+    def train_dataloader(self):
+        return DataLoader(self.cifar_train, batch_size=BATCH_SIZE, num_workers=os.cpu_count())
+
+    def val_dataloader(self):
+        return DataLoader(self.cifar_val, batch_size=BATCH_SIZE, num_workers=os.cpu_count())
+
+    def test_dataloader(self):
+        return DataLoader(self.cifar_test, batch_size=BATCH_SIZE, num_workers=os.cpu_count())

utils.py

@@ -0,0 +1,298 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+import numpy as np
+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
+
+
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+# Train data transformations
+means = [0.4914, 0.4822, 0.4465]
+stds = [0.2470, 0.2435, 0.2616]
+
+train_transforms = A.Compose(
+    [
+        A.Normalize(mean=means, std=stds, always_apply=True),
+        A.PadIfNeeded(min_height=36, min_width=36, always_apply=True),
+        A.RandomCrop(height=32, width=32, always_apply=True),
+        A.HorizontalFlip(),
+        A.CoarseDropout(max_holes=1, max_height=16, max_width=16, min_holes=1, min_height=8, min_width=8, fill_value=means),
+        ToTensorV2(),
+    ]
+)
+
+test_transforms = A.Compose(
+    [
+        A.Normalize(mean=means, std=stds, always_apply=True),
+        ToTensorV2(),
+    ]
+)
+
+
+class Cifar10SearchDataset(torchvision.datasets.CIFAR10):
+
+    def __init__(self, root="~/data", train=True, download=True, transform=None):
+        super().__init__(root=root, train=train, download=download, transform=transform)
+
+    def __getitem__(self, index):
+        image, label = self.data[index], self.targets[index]
+        if self.transform is not None:
+            transformed = self.transform(image=image)
+            image = transformed["image"]
+        return image, label
+
+def dataloader(data_path,batch_size):#,train_transforms,test_transforms):
+    trainset = Cifar10SearchDataset(root=data_path, train=True,download=True, transform=train_transforms)
+    trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,shuffle=True)
+
+    testset = Cifar10SearchDataset(root=data_path, train=False, download=True, transform=test_transforms)
+    testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,shuffle=False)
+    classes = trainset.classes
+    return trainloader, testloader, classes
+
+def plot_sample_data(dataloader):
+    batch_data, batch_label = next(iter(dataloader))
+    fig = plt.figure()
+    for i in range(12):
+        plt.subplot(3, 4, i + 1)
+        plt.tight_layout()
+        plt.imshow(torch.permute(batch_data[i], (1, 2, 0)))
+        plt.title(batch_label[i].item())
+        plt.xticks([])
+        plt.yticks([])
+
+class trainer:
+    def __init__(self,model,device,optimizer,scheduler):
+        self.model = model
+        self.device = device
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.device = device
+
+        self.train_losses = []
+        self.test_losses = []
+        self.train_acc = []
+        self.test_acc = []
+
+    def getcorrectpredcount(self,prediction, labels):
+        return prediction.argmax(dim=1).eq(labels).sum().item()
+
+    def train(self,train_loader):
+        self.model.train()
+        pbar = tqdm(train_loader)
+
+        train_loss = 0
+        correct = 0
+        processed = 0
+        criterion = nn.CrossEntropyLoss()
+
+        for batch_idx, (data, target) in enumerate(pbar):
+            data, target = data.to(self.device), target.to(self.device)
+            self.optimizer.zero_grad()
+
+            # Predict
+            pred = self.model(data)
+
+            # Calculate loss
+            loss = criterion(pred, target)
+            train_loss += loss.item()
+
+            # Backpropagation
+            loss.backward()
+            self.optimizer.step()
+
+            correct += self.getcorrectpredcount(pred, target)
+            processed += len(data)
+
+            pbar.set_description(
+                desc=f'Train: Loss={loss.item():0.4f} Batch_id={batch_idx} Accuracy={100 * correct / processed:0.2f}')
+
+        self.train_acc.append(100 * correct / processed)
+        self.train_losses.append(train_loss / len(train_loader))
+        return self.train_acc, self.train_losses
+
+    def test(self,test_loader):
+        self.model.eval()
+
+        test_loss = 0
+        correct = 0
+
+        with torch.no_grad():
+            for batch_idx, (data, target) in enumerate(test_loader):
+                data, target = data.to(self.device), target.to(self.device)
+
+                output = self.model(data)
+                test_loss += F.cross_entropy(output, target, reduction='sum').item()  # sum up batch loss
+
+                correct += self.getcorrectpredcount(output, target)
+
+        test_loss /= len(test_loader.dataset)
+        self.test_acc.append(100. * correct / len(test_loader.dataset))
+        self.test_losses.append(test_loss)
+
+        print('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format(
+            test_loss, correct, len(test_loader.dataset),
+            100. * correct / len(test_loader.dataset)))
+
+        return self.test_acc, self.test_losses
+
+    def visualize_graphs(self):
+      t = [t_items.item() for t_items in self.train_losses]
+      fig, axs = plt.subplots(2,2,figsize=(15,10))
+      axs[0, 0].plot(t)
+      axs[0, 0].set_title("Training Loss")
+      axs[1, 0].plot(self.train_acc[4000:])
+      axs[1, 0].set_title("Training Accuracy")
+      axs[0, 1].plot(self.test_losses)
+      axs[0, 1].set_title("Test Loss")
+      axs[1, 1].plot(self.test_acc)
+      axs[1, 1].set_title("Test Accuracy")
+
+    def evaluate_all_class(self,classes,test_loader):
+
+        # prepare to count predictions for each class
+        correct_pred = {classname: 0 for classname in classes}
+        total_pred = {classname: 0 for classname in classes}
+
+        # again no gradients needed
+        with torch.no_grad():
+            for data in test_loader:
+                images, labels = data
+                outputs = self.model(images)
+                _, predictions = torch.max(outputs, 1)
+                # collect the correct predictions for each class
+                for label, prediction in zip(labels, predictions):
+                    if label == prediction:
+                        correct_pred[classes[label]] += 1
+                    total_pred[classes[label]] += 1
+
+        # print accuracy for each class
+        for classname, correct_count in correct_pred.items():
+            accuracy = 100 * float(correct_count) / total_pred[classname]
+            print(f'Accuracy for class: {classname:5s} is {accuracy:.1f} %')
+        
+
+
+def evaluate_model(model, loader, device):
+    cols, rows = 4, 6
+    figure = plt.figure(figsize=(20, 20))
+    for index in range(1, cols * rows + 1):
+        k = np.random.randint(0, len(loader.dataset))  # random points from test dataset
+
+        img, label = loader.dataset[k]  # separate the image and label
+        img = img.unsqueeze(0)  # adding one dimention
+        pred = model(img.to(device))  # Prediction
+
+        figure.add_subplot(rows, cols, index)  # making the figure
+        plt.title(f"Predcited label {pred.argmax().item()}\n True Label: {label}")  # title of plot
+        plt.axis("off")  # hiding the axis
+        plt.imshow(img.squeeze(), cmap="gray")  # showing the plot
+
+    plt.show()
+
+def get_lr(optimizer):
+    """"
+    for tracking how your learning rate is changing throughout training
+    """
+    for param_group in optimizer.param_groups:
+        return param_group['lr']
+
+
+def lr_calc(model, train_loader, optimizer, criterion):
+    # model = Net().to(device)
+    # optimizer = optim.Adam(model.parameters(), lr=0.03, weight_decay=1e-4)
+    # criterion = nn.CrossEntropyLoss()
+    lr_finder = LRFinder(model, optimizer, criterion, device="cuda")
+    lr_finder.range_test(train_loader, end_lr=10, num_iter=200, step_mode="exp")
+    lr_finder.plot() # to inspect the loss-learning rate graph
+    lr_finder.reset() # to reset the model and optimizer to their initial state
+
+
+def unnormalize(img):
+    channel_means = (0.4914, 0.4822, 0.4465)
+    channel_stdevs = (0.2470, 0.2435, 0.2616)
+    img = img.numpy().astype(dtype=np.float32)
+  
+    for i in range(img.shape[0]):
+         img[i] = (img[i]*channel_stdevs[i])+channel_means[i]
+  
+    return np.transpose(img, (1,2,0))
+
+
+def plot_grad_cam_images(model, test_loader, classes, device):
+    # set model to evaluation mode
+    model.eval()
+    target_layers = [model.layer4[-1]]
+
+    # Construct the CAM object once, and then re-use it on many images:
+    cam = GradCAM(model=model, target_layers=target_layers)
+
+    misclassified_images = []
+    actual_labels = []
+    actual_targets = []
+    predicted_labels = []
+
+    with torch.no_grad():
+        for data, target in test_loader:
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            _, pred = torch.max(output, 1)
+            for i in range(len(pred)):
+                if pred[i] != target[i]:
+                    actual_targets.append(target[i])
+                    misclassified_images.append(data[i])
+                    actual_labels.append(classes[target[i]])
+                    predicted_labels.append(classes[pred[i]])
+
+    # Plot the misclassified images
+    fig = plt.figure(figsize=(12, 5))
+    for i in range(10):
+        sub = fig.add_subplot(2, 5, i+1)
+        input_tensor = misclassified_images[i].unsqueeze(dim=0)
+        targets = [ClassifierOutputTarget(actual_targets[i])]
+        grayscale_cam = cam(input_tensor=input_tensor, targets=targets)
+        grayscale_cam = grayscale_cam[0, :]
+        
+        visualization = show_cam_on_image(unnormalize(misclassified_images[i].cpu()), grayscale_cam, use_rgb=True,image_weight=0.7)
+
+        plt.imshow(visualization)
+        sub.set_title("Actual: {}, Pred: {}".format(actual_labels[i], predicted_labels[i]), color='red')
+    plt.tight_layout()
+    plt.show()
+
+def plot_misclassified_images(model, test_loader, classes, device):
+# set model to evaluation mode
+  model.eval()
+
+  misclassified_images = []
+  actual_labels = []
+  predicted_labels = []
+
+  with torch.no_grad():
+      for data, target in test_loader:
+          data, target = data.to(device), target.to(device)
+          output = model(data)
+          _, pred = torch.max(output, 1)
+          for i in range(len(pred)):
+              if pred[i] != target[i]:
+                  misclassified_images.append(data[i])
+                  actual_labels.append(classes[target[i]])
+                  predicted_labels.append(classes[pred[i]])
+
+  # Plot the misclassified images
+  fig = plt.figure(figsize=(12, 5))
+  for i in range(10):
+      sub = fig.add_subplot(2, 5, i+1)
+      npimg = unnormalize(misclassified_images[i].cpu())
+      plt.imshow(npimg, cmap='gray', interpolation='none')
+      sub.set_title("Actual: {}, Pred: {}".format(actual_labels[i], predicted_labels[i]),color='red')
+  plt.tight_layout()
+  plt.show()

visualize.py

@@ -0,0 +1,384 @@
+#!/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
+
+# 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: 'DataLoader object', 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: list[str],
+                                     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 = 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().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([])
+
+
+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: list[str], device: str, model: 'DL 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)