Adding files

README.md

@@ -1,8 +1,8 @@
 ---
-title: ERAV1 Session 12
-emoji: 😻
+title: ERA Session12
+emoji: 🔥
 colorFrom: green
-colorTo: blue
+colorTo: purple
 sdk: gradio
 sdk_version: 3.39.0
 app_file: app.py
@@ -10,4 +10,26 @@ 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.
+
+https://user-images.githubusercontent.com/23289802/258841585-4d2a75fa-3902-4839-a32a-bbfec4ef72ba.png

app.py

@@ -171,7 +171,7 @@ with gr.Blocks() as app:
                     label="Misclassified Images", info="Display misclassified images?"
                 )
                 misclassified_count = gr.Dropdown(
-                    choices=["10", "20"],
+                    choices=[str(i + 1) for i in range(20)],
                     label="Select Number of Images",
                     info="Number of Misclassified images",
                     visible=False,
@@ -188,7 +188,7 @@ with gr.Blocks() as app:
                     info="Display GradCAM images?",
                 )
                 gradcam_count = gr.Dropdown(
-                    choices=["10", "20"],
+                    choices=[str(i + 1) for i in range(20)],
                     label="Select Number of Images",
                     info="Number of GradCAM images",
                     interactive=True,
@@ -285,4 +285,4 @@ with gr.Blocks() as app:
         )
 
 
-app.launch(server_name="0.0.0.0", server_port=9998)
+app.launch()

models/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

models/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

@@ -1,228 +1,13 @@
-absl-py==1.4.0
-adbc-driver-manager==0.5.1
-adbc-driver-sqlite==0.5.1
-aiofiles==23.1.0
-aiohttp==3.8.5
-aiosignal==1.3.1
-albumentations==1.3.1
-altair==5.0.1
-annotated-types==0.5.0
-anyio==3.7.1
-argon2-cffi==21.3.0
-argon2-cffi-bindings==21.2.0
-arrow==1.2.3
-asttokens @ file:///home/conda/feedstock_root/build_artifacts/asttokens_1670263926556/work
-async-lru==2.0.4
-async-timeout==4.0.2
-attrs==23.1.0
-Babel==2.12.1
-backcall @ file:///home/conda/feedstock_root/build_artifacts/backcall_1592338393461/work
-backoff==2.2.1
-backports.functools-lru-cache @ file:///home/conda/feedstock_root/build_artifacts/backports.functools_lru_cache_1687772187254/work
-beautifulsoup4==4.12.2
-black==23.7.0
-bleach==6.0.0
-blessed==1.20.0
-cachetools==5.3.1
-certifi==2022.12.7
-cffi==1.15.1
-charset-normalizer==2.1.1
-click==8.1.6
-cloudpickle==2.2.1
-cmake==3.25.0
-connectorx==0.3.1
-contourpy==1.1.0
-croniter==1.4.1
-cycler==0.11.0
-dateutils==0.6.12
-debugpy @ file:///home/builder/ci_310/debugpy_1640789504635/work
-decorator @ file:///home/conda/feedstock_root/build_artifacts/decorator_1641555617451/work
-deepdiff==6.3.1
-defusedxml==0.7.1
-deltalake==0.10.0
-entrypoints @ file:///home/conda/feedstock_root/build_artifacts/entrypoints_1643888246732/work
-exceptiongroup==1.1.2
-executing @ file:///home/conda/feedstock_root/build_artifacts/executing_1667317341051/work
-fastapi==0.100.1
-fastjsonschema==2.18.0
-ffmpy==0.3.1
-filelock==3.12.2
-fonttools==4.41.0
-fqdn==1.5.1
-frozenlist==1.4.0
-fsspec==2023.6.0
-google-auth==2.22.0
-google-auth-oauthlib==1.0.0
-grad-cam==1.4.8
-gradio==3.39.0
-gradio_client==0.3.0
-greenlet==2.0.2
-grpcio==1.56.2
-h11==0.14.0
-httpcore==0.17.3
-httpx==0.24.1
-huggingface-hub==0.16.4
-idna==3.4
-imageio==2.31.1
-inquirer==3.1.3
-ipykernel @ file:///home/conda/feedstock_root/build_artifacts/ipykernel_1655369107642/work
-ipython @ file:///home/conda/feedstock_root/build_artifacts/ipython_1685727741709/work
-ipywidgets==8.0.7
-isoduration==20.11.0
-itsdangerous==2.1.2
-jedi @ file:///home/conda/feedstock_root/build_artifacts/jedi_1669134318875/work
-Jinja2==3.1.2
-joblib==1.3.1
-json5==0.9.14
-jsonpointer==2.4
-jsonschema==4.18.6
-jsonschema-specifications==2023.7.1
-jupyter-events==0.7.0
-jupyter-lsp==2.2.0
-jupyter_client==8.3.0
-jupyter_core @ file:///home/conda/feedstock_root/build_artifacts/jupyter_core_1686775611663/work
-jupyter_server==2.7.0
-jupyter_server_terminals==0.4.4
-jupyterlab==4.0.4
-jupyterlab-pygments==0.2.2
-jupyterlab-widgets==3.0.8
-jupyterlab_server==2.24.0
-kiwisolver==1.4.4
-lazy_loader==0.3
-lightning==2.0.6
-lightning-cloud==0.5.37
-lightning-utilities==0.9.0
-linkify-it-py==2.0.2
-lit==15.0.7
-Markdown==3.4.3
-markdown-it-py==2.2.0
-MarkupSafe==2.1.2
-matplotlib==3.7.2
-matplotlib-inline @ file:///home/conda/feedstock_root/build_artifacts/matplotlib-inline_1660814786464/work
-mdit-py-plugins==0.3.3
-mdurl==0.1.2
-mistune==3.0.1
-mpmath==1.2.1
-multidict==6.0.4
-mypy-extensions==1.0.0
-nbclient==0.8.0
-nbconvert==7.7.3
-nbformat==5.9.2
-nest-asyncio @ file:///home/conda/feedstock_root/build_artifacts/nest-asyncio_1664684991461/work
-netron==7.0.6
-networkx==3.0
-notebook_shim==0.2.3
-numpy==1.24.1
-nvidia-cublas-cu11==11.10.3.66
-nvidia-cuda-cupti-cu11==11.7.101
-nvidia-cuda-nvrtc-cu11==11.7.99
-nvidia-cuda-runtime-cu11==11.7.99
-nvidia-cudnn-cu11==8.5.0.96
-nvidia-cufft-cu11==10.9.0.58
-nvidia-curand-cu11==10.2.10.91
-nvidia-cusolver-cu11==11.4.0.1
-nvidia-cusparse-cu11==11.7.4.91
-nvidia-nccl-cu11==2.14.3
-nvidia-nvtx-cu11==11.7.91
-oauthlib==3.2.2
-opencv-python==4.8.0.74
-opencv-python-headless==4.8.0.74
-ordered-set==4.1.0
-orjson==3.9.3
-overrides==7.3.1
-packaging @ file:///home/conda/feedstock_root/build_artifacts/packaging_1681337016113/work
-pandas==2.0.3
-pandocfilters==1.5.0
-parso @ file:///home/conda/feedstock_root/build_artifacts/parso_1638334955874/work
-pathspec==0.11.2
-pexpect @ file:///home/conda/feedstock_root/build_artifacts/pexpect_1667297516076/work
-pickleshare @ file:///home/conda/feedstock_root/build_artifacts/pickleshare_1602536217715/work
-Pillow==10.0.0
-platformdirs @ file:///home/conda/feedstock_root/build_artifacts/platformdirs_1689538620473/work
-polars==0.18.8
-prometheus-client==0.17.1
-prompt-toolkit @ file:///home/conda/feedstock_root/build_artifacts/prompt-toolkit_1688565951714/work
-protobuf==4.23.4
-psutil @ file:///opt/conda/conda-bld/psutil_1656431268089/work
-ptyprocess @ file:///home/conda/feedstock_root/build_artifacts/ptyprocess_1609419310487/work/dist/ptyprocess-0.7.0-py2.py3-none-any.whl
-pure-eval @ file:///home/conda/feedstock_root/build_artifacts/pure_eval_1642875951954/work
-pyarrow==12.0.1
-pyasn1==0.5.0
-pyasn1-modules==0.3.0
-pycparser==2.21
-pydantic==2.0.3
-pydantic_core==2.3.0
-pydub==0.25.1
-Pygments @ file:///home/conda/feedstock_root/build_artifacts/pygments_1681904169130/work
-PyJWT==2.8.0
-pyparsing==3.0.9
-python-dateutil @ file:///home/conda/feedstock_root/build_artifacts/python-dateutil_1626286286081/work
-python-editor==1.0.4
-python-json-logger==2.0.7
-python-multipart==0.0.6
-pytorch-lightning==2.0.6
-pytz==2023.3
-PyWavelets==1.4.1
-PyYAML==6.0.1
-pyzmq @ file:///croot/pyzmq_1686601365461/work
-qudida==0.0.4
-readchar==4.0.5
-referencing==0.30.2
-requests==2.28.1
-requests-oauthlib==1.3.1
-rfc3339-validator==0.1.4
-rfc3986-validator==0.1.1
-rich==13.5.0
-rpds-py==0.9.2
-rsa==4.9
-ruff==0.0.280
-scikit-image==0.21.0
-scikit-learn==1.3.0
-scipy==1.11.1
-semantic-version==2.10.0
-Send2Trash==1.8.2
-six @ file:///home/conda/feedstock_root/build_artifacts/six_1620240208055/work
-sniffio==1.3.0
-soupsieve==2.4.1
-SQLAlchemy==2.0.19
-stack-data @ file:///home/conda/feedstock_root/build_artifacts/stack_data_1669632077133/work
-starlette==0.27.0
-starsessions==1.3.0
-sympy==1.11.1
-tensorboard==2.13.0
-tensorboard-data-server==0.7.1
-terminado==0.17.1
-threadpoolctl==3.2.0
-tifffile==2023.7.18
-tinycss2==1.2.1
-toml==0.10.2
-tomli==2.0.1
-toolz==0.12.0
-torch==2.0.1+cu118
-torch-lr-finder==0.2.1
-torch-tb-profiler==0.4.1
-torchaudio==2.0.2+cu118
-torchinfo==1.8.0
-torchmetrics==1.0.1
-torchvision==0.15.2+cu118
-tornado==6.3.2
-tqdm==4.65.0
-traitlets @ file:///home/conda/feedstock_root/build_artifacts/traitlets_1675110562325/work
-triton==2.0.0
-ttach==0.0.3
-typing_extensions @ file:///home/conda/feedstock_root/build_artifacts/typing_extensions_1688315532570/work
-tzdata==2023.3
-uc-micro-py==1.0.2
-uri-template==1.3.0
-urllib3==1.26.13
-uvicorn==0.23.1
-wcwidth @ file:///home/conda/feedstock_root/build_artifacts/wcwidth_1673864653149/work
-webcolors==1.13
-webencodings==0.5.1
-websocket-client==1.6.1
-websockets==11.0.3
-Werkzeug==2.3.6
-widgetsnbextension==4.0.8
-xlsx2csv==0.8.1
-XlsxWriter==3.1.2
-yarl==1.9.2
+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(),
+    ]
+)