Upload 16 files

README.md

@@ -1,7 +1,7 @@
 ---
-title: ERA Session 12
-emoji: 🐨
-colorFrom: gray
+title: ERA Session12
+emoji: 🚀
+colorFrom: green
 colorTo: purple
 sdk: gradio
 sdk_version: 3.39.0
@@ -10,4 +10,24 @@ pinned: false
 license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+### Gradio UI for CIFAR10 classification with ResNet
+
+## How to use?
+1. Select if you want visualize the misclassified images & Select the count of misclassified images.
+2. Select if you want to visualize the GradCAM images & Also select count of Gradcam images, Model layer and Opacity of the resulting image.
+3. Click on the upload button to upload the local image to be used for prediction and select the image for prediction.
+4. If you want use one of the sample images, please pick one from the list of 10 sample images.
+5. Select the top n classes for which you want see the model performance.
+6. Click on the Run button
+7. On the right side of the interface, the top view displays the selected number of misclassified images.
+8. The second view displays the GradCAM output.
+9. And Final view displays the top n predicitons for the given image.
+
+## Components Used:
+1. `gr.Dropdown` : Used for selecting the number of images for Misclassified & GradCAM output and also for the top n classes to be displayed.
+2. `gr.Checkbox` : Used for boolean inputs like if user wants to visualize Misclassified or if they want to visualize gradCAM images.
+3. `gr.Slider` : Used to select the opacity paramter to be used with GradCAM viaualization.
+4. `gr.Gallery`: Used to display a numebr of images, used for displaying input images and output images.
+5. `gr.UploadButton`: A generic file uplaod button, used for picking and uploading local image file for prediction.
+6. `gr.Button`: Used for calling the main prediction module.
+7. `gr.Label`: Used for displaying the top n classification results.

app.py

@@ -0,0 +1,288 @@
+import gradio as gr
+import random
+import numpy as np
+from PIL import Image
+import torch
+import torchvision
+
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+
+from models.resnet_lightning import ResNet
+from utils.data import CIFARDataModule
+from utils.transforms import test_transform
+from utils.common import get_misclassified_data
+
+inv_normalize = torchvision.transforms.Normalize(
+    mean=[-0.50 / 0.23, -0.50 / 0.23, -0.50 / 0.23], std=[1 / 0.23, 1 / 0.23, 1 / 0.23]
+)
+
+datamodule = CIFARDataModule()
+datamodule.setup()
+classes = datamodule.train_dataset.classes
+
+model = ResNet.load_from_checkpoint("model.ckpt")
+model = model.to("cpu")
+
+prediction_image = None
+
+
+def upload_file(files):
+    file_paths = [file.name for file in files]
+    return file_paths
+
+
+def read_image(path):
+    img = Image.open(path)
+    img.load()
+    data = np.asarray(img, dtype="uint8")
+    return data
+
+
+def sample_images():
+    images = []
+    length = len(datamodule.test_dataset)
+    classes = datamodule.train_dataset.classes
+    for i in range(10):
+        idx = random.randint(0, length - 1)
+        image, label = datamodule.test_dataset[idx]
+        image = inv_normalize(image).permute(1, 2, 0).numpy()
+        images.append((image, classes[label]))
+    return images
+
+
+def get_misclassified_images(misclassified_count):
+    misclassified_images = []
+    misclassified_data = get_misclassified_data(
+        model=model,
+        device="cpu",
+        test_loader=datamodule.test_dataloader(),
+        count=misclassified_count,
+    )
+    for i in range(misclassified_count):
+        img = misclassified_data[i][0].squeeze().to("cpu")
+        img = inv_normalize(img)
+        img = np.transpose(img.numpy(), (1, 2, 0))
+        label = f"Label: {classes[misclassified_data[i][1].item()]} | Prediction: {classes[misclassified_data[i][2].item()]}"
+        misclassified_images.append((img, label))
+    return misclassified_images
+
+
+def get_gradcam_images(gradcam_layer, gradcam_count, gradcam_opacity):
+    gradcam_images = []
+    if gradcam_layer == "Layer1":
+        target_layers = [model.layer1[-1]]
+    elif gradcam_layer == "Layer2":
+        target_layers = [model.layer2[-1]]
+    else:
+        target_layers = [model.layer3[-1]]
+
+    cam = GradCAM(model=model, target_layers=target_layers, use_cuda=False)
+    data = get_misclassified_data(
+        model=model,
+        device="cpu",
+        test_loader=datamodule.test_dataloader(),
+        count=gradcam_count,
+    )
+    for i in range(gradcam_count):
+        input_tensor = data[i][0]
+
+        # Get the activations of the layer for the images
+        grayscale_cam = cam(input_tensor=input_tensor, targets=None)
+        grayscale_cam = grayscale_cam[0, :]
+
+        # Get back the original image
+        img = input_tensor.squeeze(0).to("cpu")
+        if inv_normalize is not None:
+            img = inv_normalize(img)
+        rgb_img = np.transpose(img, (1, 2, 0))
+        rgb_img = rgb_img.numpy()
+
+        # Mix the activations on the original image
+        visualization = show_cam_on_image(
+            rgb_img, grayscale_cam, use_rgb=True, image_weight=gradcam_opacity
+        )
+        label = f"Label: {classes[data[i][1].item()]} | Prediction: {classes[data[i][2].item()]}"
+        gradcam_images.append((visualization, label))
+    return gradcam_images
+
+
+def show_hide_misclassified(status):
+    if not status:
+        return {misclassified_count: gr.update(visible=False)}
+    return {misclassified_count: gr.update(visible=True)}
+
+
+def show_hide_gradcam(status):
+    if not status:
+        return [gr.update(visible=False) for i in range(3)]
+    return [gr.update(visible=True) for i in range(3)]
+
+
+def set_prediction_image(evt: gr.SelectData, gallery):
+    global prediction_image
+    if isinstance(gallery[evt.index], dict):
+        prediction_image = gallery[evt.index]["name"]
+    else:
+        prediction_image = gallery[evt.index][0]["name"]
+
+
+def predict(
+    is_misclassified,
+    misclassified_count,
+    is_gradcam,
+    gradcam_count,
+    gradcam_layer,
+    gradcam_opacity,
+    num_classes,
+):
+    misclassified_images = None
+    if is_misclassified:
+        misclassified_images = get_misclassified_images(int(misclassified_count))
+
+    gradcam_images = None
+    if is_gradcam:
+        gradcam_images = get_gradcam_images(
+            gradcam_layer, int(gradcam_count), gradcam_opacity
+        )
+
+    img = read_image(prediction_image)
+    image_transformed = test_transform(image=img)["image"]
+    output = model(image_transformed.unsqueeze(0))
+    preds = torch.softmax(output, dim=1).squeeze().detach().numpy()
+    indices = (
+        output.argsort(descending=True).squeeze().detach().numpy()[: int(num_classes)]
+    )
+    predictions = {classes[i]: round(float(preds[i]), 2) for i in indices}
+
+    return {
+        miscalssfied_output: gr.update(value=misclassified_images),
+        gradcam_output: gr.update(value=gradcam_images),
+        prediction_label: gr.update(value=predictions),
+    }
+
+
+with gr.Blocks() as app:
+    gr.Markdown("## ERA Session12 - CIFAR10 Classification with ResNet")
+    with gr.Row():
+        with gr.Column():
+            with gr.Box():
+                is_misclassified = gr.Checkbox(
+                    label="Misclassified Images", info="Display misclassified images?"
+                )
+                misclassified_count = gr.Dropdown(
+                    choices=[str(i + 1) for i in range(20)],
+                    label="Select Number of Images",
+                    info="Number of Misclassified images",
+                    visible=False,
+                    interactive=True,
+                )
+                is_misclassified.input(
+                    show_hide_misclassified,
+                    inputs=[is_misclassified],
+                    outputs=[misclassified_count],
+                )
+            with gr.Box():
+                is_gradcam = gr.Checkbox(
+                    label="GradCAM Images",
+                    info="Display GradCAM images?",
+                )
+                gradcam_count = gr.Dropdown(
+                    choices=[str(i + 1) for i in range(20)],
+                    label="Select Number of Images",
+                    info="Number of GradCAM images",
+                    interactive=True,
+                    visible=False,
+                )
+                gradcam_layer = gr.Dropdown(
+                    choices=["Layer1", "Layer2", "Layer3"],
+                    label="Select the layer",
+                    info="Please select the layer for which the GradCAM is required",
+                    interactive=True,
+                    visible=False,
+                )
+                gradcam_opacity = gr.Slider(
+                    minimum=0,
+                    maximum=1,
+                    value=0.6,
+                    label="Opacity",
+                    info="Opacity of GradCAM output",
+                    interactive=True,
+                    visible=False,
+                )
+
+                is_gradcam.input(
+                    show_hide_gradcam,
+                    inputs=[is_gradcam],
+                    outputs=[gradcam_count, gradcam_layer, gradcam_opacity],
+                )
+            with gr.Box():
+                # file_output = gr.File(file_types=["image"])
+                with gr.Group():
+                    upload_gallery = gr.Gallery(
+                        value=None,
+                        label="Uploaded images",
+                        show_label=False,
+                        elem_id="gallery_upload",
+                        columns=5,
+                        rows=2,
+                        height="auto",
+                        object_fit="contain",
+                    )
+                    upload_button = gr.UploadButton(
+                        "Click to Upload images",
+                        file_types=["image"],
+                        file_count="multiple",
+                    )
+                    upload_button.upload(upload_file, upload_button, upload_gallery)
+
+                with gr.Group():
+                    sample_gallery = gr.Gallery(
+                        value=sample_images,
+                        label="Sample images",
+                        show_label=True,
+                        elem_id="gallery_sample",
+                        columns=5,
+                        rows=2,
+                        height="auto",
+                        object_fit="contain",
+                    )
+
+                upload_gallery.select(set_prediction_image, inputs=[upload_gallery])
+                sample_gallery.select(set_prediction_image, inputs=[sample_gallery])
+
+            with gr.Box():
+                num_classes = gr.Dropdown(
+                    choices=[str(i + 1) for i in range(10)],
+                    label="Select Number of Top Classes",
+                    info="Number of Top target classes to be shown",
+                )
+            run_btn = gr.Button()
+        with gr.Column():
+            with gr.Box():
+                miscalssfied_output = gr.Gallery(
+                    value=None, label="Misclassified Images", show_label=True
+                )
+            with gr.Box():
+                gradcam_output = gr.Gallery(
+                    value=None, label="GradCAM Images", show_label=True
+                )
+            with gr.Box():
+                prediction_label = gr.Label(value=None, label="Predictions")
+
+        run_btn.click(
+            predict,
+            inputs=[
+                is_misclassified,
+                misclassified_count,
+                is_gradcam,
+                gradcam_count,
+                gradcam_layer,
+                gradcam_opacity,
+                num_classes,
+            ],
+            outputs=[miscalssfied_output, gradcam_output, prediction_label],
+        )
+
+
+app.launch()

best_model.pth

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

config.toml

@@ -0,0 +1,13 @@
+[data]
+batch_size = 512
+shuffle = true
+num_workers = 4
+
+[training]
+epochs = 20
+batch_size = 512
+optimizer = "adam"
+criterion = "crossentropy"
+lr = 0.003
+weight_decay = 1e-4
+lrfinder = { numiter = 600, endlr = 10, startlr = 1e-2 }

gitattributes.txt

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

gitignore.txt

@@ -0,0 +1,4 @@
+lightning_logs
+data
+.ipynb_checkpoints
+__pycache__/

model/custom_resnet.py

@@ -0,0 +1,112 @@
+import torch
+import torch.nn as nn
+
+
+class ResBlock(nn.Module):
+    def __init__(self, channels):
+        super(ResBlock, self).__init__()
+
+        self.resblock = nn.Sequential(
+            nn.Conv2d(
+                in_channels=channels,
+                out_channels=channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(channels),
+            nn.ReLU(),
+            nn.Conv2d(
+                in_channels=channels,
+                out_channels=channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(channels),
+            nn.ReLU(),
+        )
+
+    def forward(self, x):
+        return x + self.resblock(x)
+
+
+class CustomResnet(nn.Module):
+    def __init__(self):
+        super(CustomResnet, self).__init__()
+
+        self.prep = nn.Sequential(
+            nn.Conv2d(
+                in_channels=3,
+                out_channels=64,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+        )
+
+        self.layer1 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=64,
+                out_channels=128,
+                kernel_size=3,
+                padding=1,
+                stride=1,
+                bias=False,
+            ),
+            nn.MaxPool2d(kernel_size=2),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            ResBlock(channels=128),
+        )
+
+        self.layer2 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=128,
+                out_channels=256,
+                kernel_size=3,
+                padding=1,
+                stride=1,
+                bias=False,
+            ),
+            nn.MaxPool2d(kernel_size=2),
+            nn.BatchNorm2d(256),
+            nn.ReLU(),
+        )
+
+        self.layer3 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=256,
+                out_channels=512,
+                kernel_size=3,
+                padding=1,
+                stride=1,
+                bias=False,
+            ),
+            nn.MaxPool2d(kernel_size=2),
+            nn.BatchNorm2d(512),
+            nn.ReLU(),
+            ResBlock(channels=512),
+        )
+
+        self.pool = nn.MaxPool2d(kernel_size=4)
+
+        self.fc = nn.Linear(in_features=512, out_features=10, bias=False)
+
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x):
+        x = self.prep(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.pool(x)
+        x = x.view(-1, 512)
+        x = self.fc(x)
+        # x = self.softmax(x)
+        return x

model/resnet_lightning.py

@@ -0,0 +1,173 @@
+import torch
+import torch.nn as nn
+import lightning as L
+from torchmetrics import Accuracy
+from typing import Any
+
+from utils.common import one_cycle_lr
+
+class ResidualBlock(L.LightningModule):
+    def __init__(self, channels):
+        super(ResidualBlock, self).__init__()
+
+        self.residual_block = nn.Sequential(
+            nn.Conv2d(
+                in_channels=channels,
+                out_channels=channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(channels),
+            nn.ReLU(),
+            nn.Conv2d(
+                in_channels=channels,
+                out_channels=channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(channels),
+            nn.ReLU(),
+        )
+
+    def forward(self, x):
+        return x + self.residual_block(x)
+
+class ResNet(L.LightningModule):
+    def __init__(
+        self, batch_size=512, shuffle=True, num_workers=4, learning_rate=0.003, scheduler_steps=None, maxlr=None, epochs=None
+    ):
+        super(ResNet, self).__init__()
+        self.data_dir = "./data"
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.num_workers = num_workers
+        self.learning_rate = learning_rate
+        self.scheduler_steps = scheduler_steps
+        self.maxlr = maxlr if maxlr is not None else learning_rate
+        self.epochs = epochs
+
+        self.prep = nn.Sequential(
+            nn.Conv2d(
+                in_channels=3,
+                out_channels=64,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+        )
+
+        self.layer1 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=64,
+                out_channels=128,
+                kernel_size=3,
+                padding=1,
+                stride=1,
+                bias=False,
+            ),
+            nn.MaxPool2d(kernel_size=2),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            ResidualBlock(channels=128),
+        )
+
+        self.layer2 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=128,
+                out_channels=256,
+                kernel_size=3,
+                padding=1,
+                stride=1,
+                bias=False,
+            ),
+            nn.MaxPool2d(kernel_size=2),
+            nn.BatchNorm2d(256),
+            nn.ReLU(),
+        )
+
+        self.layer3 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=256,
+                out_channels=512,
+                kernel_size=3,
+                padding=1,
+                stride=1,
+                bias=False,
+            ),
+            nn.MaxPool2d(kernel_size=2),
+            nn.BatchNorm2d(512),
+            nn.ReLU(),
+            ResidualBlock(channels=512),
+        )
+
+        self.pool = nn.MaxPool2d(kernel_size=4)
+
+        self.fc = nn.Linear(in_features=512, out_features=10, bias=False)
+
+        self.softmax = nn.Softmax(dim=-1)
+
+        self.accuracy = Accuracy(task="multiclass", num_classes=10)
+
+    def forward(self, x):
+        x = self.prep(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.pool(x)
+        x = x.view(-1, 512)
+        x = self.fc(x)
+        # x = self.softmax(x)
+        return x
+
+    def configure_optimizers(self) -> Any:
+        optimizer = torch.optim.Adam(
+            self.parameters(), lr=self.learning_rate, weight_decay=1e-4
+        )
+        scheduler = one_cycle_lr(
+    optimizer=optimizer, maxlr=self.maxlr, steps=self.scheduler_steps, epochs=self.epochs
+)
+        return {"optimizer": optimizer,
+                "lr_scheduler": {"scheduler": scheduler,
+                                 "interval": "step"}}
+
+    def training_step(self, batch, batch_idx):
+        X, y = batch
+        y_pred = self(X)
+        loss = nn.CrossEntropyLoss()(y_pred, y)
+
+        preds = torch.argmax(y_pred, dim=1)
+
+        accuracy = self.accuracy(preds, y)
+
+        self.log_dict({"train_loss": loss, "train_acc": accuracy}, prog_bar=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        X, y = batch
+        y_pred = self(X)
+        loss = nn.CrossEntropyLoss(reduction="sum")(y_pred, y)
+
+        preds = torch.argmax(y_pred, dim=1)
+
+        accuracy = self.accuracy(preds, y)
+
+        self.log_dict({"val_loss": loss, "val_acc": accuracy}, prog_bar=True)
+
+        return loss
+
+    def test_step(self, batch, batch_idx):
+        X, y = batch
+        y_pred = self(X)
+        loss = nn.CrossEntropyLoss(reduction="sum")(y_pred, y)
+        preds = torch.argmax(y_pred, dim=1)
+
+        accuracy = self.accuracy(preds, y)
+
+        self.log_dict({"test_loss": loss, "test_acc": accuracy}, prog_bar=True)

requirements.txt

@@ -0,0 +1,13 @@
+numpy
+pandas
+matplotlib
+torch
+torchvision
+lightning
+gradio
+grad-cam
+torchinfo
+torch_lr_finder
+pydantic
+tqdm
+albumentations

utils/common.py

@@ -0,0 +1,185 @@
+import numpy as np
+import random
+import matplotlib.pyplot as plt
+
+import torch
+import torchvision
+from torchinfo import summary
+from torch_lr_finder import LRFinder
+
+
+def find_lr(model, optimizer, criterion, device, trainloader, numiter, startlr, endlr):
+    lr_finder = LRFinder(
+        model=model, optimizer=optimizer, criterion=criterion, device=device
+    )
+
+    lr_finder.range_test(
+        train_loader=trainloader,
+        start_lr=startlr,
+        end_lr=endlr,
+        num_iter=numiter,
+        step_mode="exp",
+    )
+
+    lr_finder.plot()
+
+    lr_finder.reset()
+
+
+def one_cycle_lr(optimizer, maxlr, steps, epochs):
+    scheduler = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer=optimizer,
+        max_lr=maxlr,
+        steps_per_epoch=steps,
+        epochs=epochs,
+        pct_start=5 / epochs,
+        div_factor=100,
+        three_phase=False,
+        final_div_factor=100,
+        anneal_strategy="linear",
+    )
+    return scheduler
+
+
+def show_random_images_for_each_class(train_data, num_images_per_class=16):
+    for c, cls in enumerate(train_data.classes):
+        rand_targets = random.sample(
+            [n for n, x in enumerate(train_data.targets) if x == c],
+            k=num_images_per_class,
+        )
+        show_img_grid(np.transpose(train_data.data[rand_targets], axes=(0, 3, 1, 2)))
+        plt.title(cls)
+
+
+def show_img_grid(data):
+    try:
+        grid_img = torchvision.utils.make_grid(data.cpu().detach())
+    except:
+        data = torch.from_numpy(data)
+        grid_img = torchvision.utils.make_grid(data)
+
+    plt.figure(figsize=(10, 10))
+    plt.imshow(grid_img.permute(1, 2, 0))
+
+
+def show_random_images(data_loader):
+    data, target = next(iter(data_loader))
+    show_img_grid(data)
+
+
+def show_model_summary(model, batch_size):
+    summary(
+        model=model,
+        input_size=(batch_size, 3, 32, 32),
+        col_names=["input_size", "output_size", "num_params", "kernel_size"],
+        verbose=1,
+    )
+
+
+def lossacc_plots(results):
+    plt.plot(results["epoch"], results["trainloss"])
+    plt.plot(results["epoch"], results["testloss"])
+    plt.legend(["Train Loss", "Validation Loss"])
+    plt.xlabel("Epochs")
+    plt.ylabel("Loss")
+    plt.title("Loss vs Epochs")
+    plt.show()
+
+    plt.plot(results["epoch"], results["trainacc"])
+    plt.plot(results["epoch"], results["testacc"])
+    plt.legend(["Train Acc", "Validation Acc"])
+    plt.xlabel("Epochs")
+    plt.ylabel("Accuracy")
+    plt.title("Accuracy vs Epochs")
+    plt.show()
+
+
+def lr_plots(results, length):
+    plt.plot(range(length), results["lr"])
+    plt.xlabel("Epochs")
+    plt.ylabel("Learning Rate")
+    plt.title("Learning Rate vs Epochs")
+    plt.show()
+
+
+def get_misclassified(model, testloader, device, mis_count=10):
+    misimgs, mistgts, mispreds = [], [], []
+    with torch.no_grad():
+        for data, target in testloader:
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            pred = output.argmax(dim=1, keepdim=True)
+            misclassified = torch.argwhere(pred.squeeze() != target).squeeze()
+            for idx in misclassified:
+                if len(misimgs) >= mis_count:
+                    break
+                misimgs.append(data[idx])
+                mistgts.append(target[idx])
+                mispreds.append(pred[idx].squeeze())
+    return misimgs, mistgts, mispreds
+
+
+# def plot_misclassified(misimgs, mistgts, mispreds, classes):
+#     fig, axes = plt.subplots(len(misimgs) // 2, 2)
+#     fig.tight_layout()
+#     for ax, img, tgt, pred in zip(axes.ravel(), misimgs, mistgts, mispreds):
+#         ax.imshow((img / img.max()).permute(1, 2, 0).cpu())
+#         ax.set_title(f"{classes[tgt]} | {classes[pred]}")
+#         ax.grid(False)
+#         ax.set_axis_off()
+#     plt.show()
+
+def get_misclassified_data(model, device, test_loader, count):
+    """
+    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))
+            if len(misclassified_data) >= count:
+                        break
+            
+    return misclassified_data[:count]
+
+def plot_misclassified(data, classes, size=(10, 10), rows=2, cols=5, inv_normalize=None):
+    fig = plt.figure(figsize=size)
+    number_of_samples = len(data)
+    for i in range(number_of_samples):
+        plt.subplot(rows, cols, i + 1)
+        img = data[i][0].squeeze().to('cpu')
+        if inv_normalize is not None:
+            img = inv_normalize(img)
+        plt.imshow(np.transpose(img, (1, 2, 0)))
+        plt.title(f"Label: {classes[data[i][1].item()]} \n Prediction: {classes[data[i][2].item()]}")
+        plt.xticks([])
+        plt.yticks([])
+

utils/config.py

@@ -0,0 +1,36 @@
+import toml
+from pydantic import BaseModel
+
+TOML_PATH = "config.toml"
+
+
+class Data(BaseModel):
+    batch_size: int = 512
+    shuffle: bool = True
+    num_workers: int = 4
+
+
+class LRFinder(BaseModel):
+    numiter: int = 600
+    endlr: float = 10
+    startlr: float = 1e-2
+
+
+class Training(BaseModel):
+    epochs: int = 20
+    optimizer: str = "adam"
+    criterion: str = "crossentropy"
+    lr: float = 0.003
+    weight_decay: float = 1e-4
+    lrfinder: LRFinder
+
+
+class Config(BaseModel):
+    data: Data
+    training: Training
+
+
+with open(TOML_PATH) as f:
+    toml_config = toml.load(f)
+
+config = Config(**toml_config)

utils/data.py

@@ -0,0 +1,68 @@
+import torchvision
+import lightning as L
+from torch.utils.data import DataLoader
+from utils.transforms import train_transform, test_transform
+
+
+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
+
+
+class CIFARDataModule(L.LightningDataModule):
+    def __init__(
+        self, data_dir="data", batch_size=512, shuffle=True, num_workers=4
+    ) -> None:
+        super().__init__()
+        self.data_dir = data_dir
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.num_workers = num_workers
+
+    def prepare_data(self) -> None:
+        pass
+
+    def setup(self, stage=None):
+        self.train_dataset = Cifar10SearchDataset(
+            root=self.data_dir, train=True, transform=train_transform
+        )
+
+        self.val_dataset = Cifar10SearchDataset(
+            root=self.data_dir, train=False, transform=test_transform
+        )
+
+        self.test_dataset = Cifar10SearchDataset(
+            root=self.data_dir, train=False, transform=test_transform
+        )
+
+    def train_dataloader(self):
+        return DataLoader(
+            dataset=self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=self.shuffle,
+            num_workers=self.num_workers,
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            dataset=self.val_dataset,
+            batch_size=self.batch_size,
+            shuffle=self.shuffle,
+            num_workers=self.num_workers,
+        )
+
+    def test_dataloader(self):
+        return DataLoader(
+            dataset=self.test_dataset,
+            batch_size=self.batch_size,
+            shuffle=self.shuffle,
+            num_workers=self.num_workers,
+        )

utils/gradcam.py

@@ -0,0 +1,67 @@
+import numpy as np
+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 matplotlib.pyplot as plt
+
+
+def generate_gradcam(model, target_layers, images, labels, rgb_imgs):
+    results = []
+    cam = GradCAM(model=model, target_layers=target_layers, use_cuda=True)
+
+    for image, label, np_image in zip(images, labels, rgb_imgs):
+        targets = [ClassifierOutputTarget(label.item())]
+
+        # You can also pass aug_smooth=True and eigen_smooth=True, to apply smoothing.
+        grayscale_cam = cam(
+            input_tensor=image.unsqueeze(0), targets=targets, aug_smooth=True
+        )
+
+        # In this example grayscale_cam has only one image in the batch:
+        grayscale_cam = grayscale_cam[0, :]
+        visualization = show_cam_on_image(
+            np_image / np_image.max(), grayscale_cam, use_rgb=True
+        )
+        results.append(visualization)
+    return results
+
+
+def visualize_gradcam(misimgs, mistgts, mispreds, classes):
+    fig, axes = plt.subplots(len(misimgs) // 2, 2)
+    fig.tight_layout()
+    for ax, img, tgt, pred in zip(axes.ravel(), misimgs, mistgts, mispreds):
+        ax.imshow(img)
+        ax.set_title(f"{classes[tgt]} | {classes[pred]}")
+        ax.grid(False)
+        ax.set_axis_off()
+    plt.show()
+
+def plot_gradcam(model, data, classes, target_layers, number_of_samples, inv_normalize=None, targets=None, transparency = 0.60, figsize=(10,10), rows=2, cols=5):
+    
+    fig = plt.figure(figsize=figsize)
+    
+    cam = GradCAM(model=model, target_layers=target_layers, use_cuda=True)
+    for i in range(number_of_samples):
+            plt.subplot(rows, cols, 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')
+            if inv_normalize is not None:
+                img = inv_normalize(img)
+            rgb_img = np.transpose(img, (1, 2, 0))
+            rgb_img = rgb_img.numpy()
+    
+            # 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(f"Label: {classes[data[i][1].item()]} \n Prediction: {classes[data[i][2].item()]}")
+            plt.xticks([])
+            plt.yticks([])

utils/training.py

@@ -0,0 +1,90 @@
+from tqdm import tqdm
+import torch
+import torch.nn.functional as F
+
+
+def train(
+    model,
+    device,
+    train_loader,
+    optimizer,
+    criterion,
+    scheduler,
+    L1=False,
+    l1_lambda=0.01,
+):
+    model.train()
+    pbar = tqdm(train_loader)
+
+    train_losses = []
+    train_acc = []
+    lrs = []
+
+    correct = 0
+    processed = 0
+    train_loss = 0
+
+    for batch_idx, (data, target) in enumerate(pbar):
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        y_pred = model(data)
+
+        # Calculate loss
+        loss = criterion(y_pred, target)
+        if L1:
+            l1_loss = 0
+            for p in model.parameters():
+                l1_loss = l1_loss + p.abs().sum()
+            loss = loss + l1_lambda * l1_loss
+        else:
+            loss = loss
+
+        train_loss += loss.item()
+        train_losses.append(loss.item())
+
+        # Backpropagation
+        loss.backward()
+        optimizer.step()
+        scheduler.step()
+
+        # Update pbar-tqdm
+        pred = y_pred.argmax(
+            dim=1, keepdim=True
+        )  # get the index of the max log-probability
+        correct += pred.eq(target.view_as(pred)).sum().item()
+        processed += len(data)
+
+        pbar.set_description(
+            desc=f"Loss={loss.item():0.2f} Accuracy={100*correct/processed:0.2f}"
+        )
+        train_acc.append(100 * correct / processed)
+        lrs.append(scheduler.get_last_lr())
+
+    return train_losses, train_acc, lrs
+
+
+def test(model, device, criterion, test_loader):
+    model.eval()
+    test_loss = 0
+    correct = 0
+    with torch.no_grad():
+        for data, target in test_loader:
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            test_loss += F.cross_entropy(output, target, reduction="sum").item()
+            pred = output.argmax(dim=1, keepdim=True)
+            correct += pred.eq(target.view_as(pred)).sum().item()
+
+    test_loss /= len(test_loader.dataset)
+
+    print(
+        "\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n".format(
+            test_loss,
+            correct,
+            len(test_loader.dataset),
+            100.0 * correct / len(test_loader.dataset),
+        )
+    )
+    test_acc = 100.0 * correct / len(test_loader.dataset)
+
+    return test_loss, test_acc

utils/transforms.py

@@ -0,0 +1,31 @@
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+train_transform = A.Compose(
+    [
+        A.PadIfNeeded(min_height=40, min_width=40, always_apply=True),
+        A.RandomCrop(height=32, width=32, always_apply=True),
+        A.HorizontalFlip(),
+        A.CoarseDropout(
+            min_holes=1,
+            max_holes=1,
+            min_height=8,
+            min_width=8,
+            max_height=8,
+            max_width=8,
+            fill_value=[0.49139968*255, 0.48215827*255 ,0.44653124*255],  # type: ignore
+            p=0.5,
+        ),
+        A.Normalize((0.49139968, 0.48215827, 0.44653124),
+                    (0.24703233, 0.24348505, 0.26158768)),
+        ToTensorV2(),
+    ]
+)
+
+test_transform = A.Compose(
+    [
+        A.Normalize((0.49139968, 0.48215827, 0.44653124),
+                    (0.24703233, 0.24348505, 0.26158768)),
+        ToTensorV2(),
+    ]
+)