Upload 17 files

Store/model.pth

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

Utilities/config.py

@@ -0,0 +1,65 @@
+import torch.nn.functional as F
+
+
+#SEED
+
+SEED = 1
+
+#DATASET
+
+CLASSES = (
+    "Airplane",
+    "Automobile",
+    "Bird",
+    "Cat",
+    "Deer",
+    "Dog",
+    "Frog",
+    "Horse",
+    "Ship",
+    "Truck"
+)
+
+SHUFFLE = True
+DATA_DIR = "../data"
+NUM_WORKERS = 4
+PIN_MEMORY = True
+
+# TRAINING HP
+
+CRITERION = F.cross_entropy
+INPUT_SIZE = (3, 32, 32)
+NUM_CLASSES = 10
+LEARNING_RATE = 0.001
+WEIGHT_DECAY = 1e-4
+BATCH_SIZE = 512
+NUM_EPOCHS = 24
+DROPOUT_PERCENTAGE = 0.05
+LAYER_NORM = "bn"
+
+# OPTIMIZER & SCHEDULAR
+
+LRFINDER_END_LR = 0.1
+LRFINDER_NUM_ITERATIONS = 50
+LRFINDER_STEP_MODE = "exp"
+
+OCLR_DIV_FACTOR = 100
+OCLR_FINAL_DIV_FACTOR = 100
+OCLR_THREE_PHASE = False
+OCLR_ANNEAL_STRATEGY = "linear"
+
+# COMPUTE RELATED
+
+ACCELERATOR = "cpu"
+PRECISION = 32
+
+# STORAGE
+
+TRAINING_STAT_STORE = "Store/training_stats.csv"
+MODEL_SAVE_PATH = "Store/model.pth"
+PRED_STORE_PATH = "Store/pred_store.pth"
+EXAMPLE_IMG_PATH = "Store/examples/"
+
+# VISULIZATION
+
+NORM_CONF_MAT = True

Utilities/model.py

@@ -0,0 +1,305 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import pytorch_lightning as pl
+import seaborn as sns
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torchmetrics
+from torch.optim.lr_scheduler import OneCycleLR
+from torch_lr_finder import LRFinder
+
+from . import config    # Custom config file
+from .visualize import plot_incorrect_preds
+
+
+class Net(pl.LightningModule):
+    def __init__(
+            self,
+            num_classes=10,
+            dropout_percentage=0,
+            norm='bn',
+            num_groups=2,
+            criterion=F.cross_entropy,
+            learning_rate=0.001,
+            weight_decay=0.0
+        ):
+        super(Net, self).__init__()
+
+        # Define norm
+        if norm == 'bn':
+            self.norm = nn.BatchNorm2d
+        elif norm == 'gn':
+            self.norm = lambda in_dim: nn.GroupNorm(
+                num_groups=num_groups, num_channels=in_dim
+            )
+        elif norm == 'ln':
+            self.norm = lambda in_dim: nn.GroupNorm(
+                num_groups=in_dim, num_channels=in_dim
+            )
+        
+        #define loss
+        self.criterion = criterion
+
+        #define metrics
+        self.accuracy = torchmetrics.Accuracy(
+            task='multiclass', num_classes=num_classes
+        )
+        self.confusion_matrix = torchmetrics.ConfusionMatrix(
+            task='multiclass', num_classes=num_classes
+        )
+        
+        #define the optimizer hyperparameters
+        self.learning_rate = learning_rate
+        self.weight_decay = weight_decay
+
+        #prediction storage
+        self.pred_store = {
+            "test_preds": torch.tensor([]),
+            "test_labels": torch.tensor([]),
+            "test_incorrect": [] #?
+        }
+        self.log_store = {          # not used at all
+            "train_loss_epoch": [],
+            "train_acc_epoch": [],
+            "val_loss_epoch": [],
+            "val_acc_epoch": [],
+            "test_loss_epoch": [],  # not used
+            "test_acc_epoch": [],   # not used
+        }
+
+        # Define the network architecture
+        self.prep_layer = nn.Sequential(
+            nn.Conv2d(3, 64, kernel_size=3, padding=1),  # 32x32x3 | 1 -> 32x32x64 | 3
+            self.norm(64),
+            nn.ReLU(),
+            nn.Dropout(dropout_percentage),
+        )
+
+        self.l1 = nn.Sequential(
+            nn.Conv2d(64, 128, kernel_size=3, padding=1),  # 32x32x128 | 5
+            nn.MaxPool2d(2, 2),  # 16x16x128 | 6
+            self.norm(128),
+            nn.ReLU(),
+            nn.Dropout(dropout_percentage),
+        )
+        self.l1res = nn.Sequential(
+            nn.Conv2d(128, 128, kernel_size=3, padding=1),  # 16x16x128 | 10
+            self.norm(128),
+            nn.ReLU(),
+            nn.Dropout(dropout_percentage),
+            nn.Conv2d(128, 128, kernel_size=3, padding=1),  # 16x16x128 | 14
+            self.norm(128),
+            nn.ReLU(),
+            nn.Dropout(dropout_percentage),
+        )
+        self.l2 = nn.Sequential(
+            nn.Conv2d(128, 256, kernel_size=3, padding=1),  # 16x16x256 | 18
+            nn.MaxPool2d(2, 2),  # 8x8x256 | 19
+            self.norm(256),
+            nn.ReLU(),
+            nn.Dropout(dropout_percentage),
+        )
+        self.l3 = nn.Sequential(
+            nn.Conv2d(256, 512, kernel_size=3, padding=1),  # 8x8x512 | 27
+            nn.MaxPool2d(2, 2),  # 4x4x512 | 28
+            self.norm(512),
+            nn.ReLU(),
+            nn.Dropout(dropout_percentage),
+        )
+        self.l3res = nn.Sequential(
+            nn.Conv2d(512, 512, kernel_size=3, padding=1),  # 4x4x512 | 36
+            self.norm(512),
+            nn.ReLU(),
+            nn.Dropout(dropout_percentage),
+            nn.Conv2d(512, 512, kernel_size=3, padding=1),  # 4x4x512 | 44
+            self.norm(512),
+            nn.ReLU(),
+            nn.Dropout(dropout_percentage),
+        )
+        self.maxpool = nn.MaxPool2d(4, 4)
+
+        # Classifier
+        self.linear = nn.Linear(512, 10)
+
+    def forward(self, x):
+        x = self.prep_layer(x)
+        x = self.l1(x)
+        x = x + self.l1res(x)
+        x = self.l2(x)
+        x = self.l3(x)
+        x = x + self.l3res(x)
+        x = self.maxpool(x)
+        x = x.view(-1, 512)
+        x = self.linear(x)
+        return F.log_softmax(x, dim=1)
+    
+    def training_step(self, batch, batch_idx):
+        data, target = batch
+
+        #forward pass
+        pred = self.forward(data)
+
+        #calculate loss
+        loss = self.criterion(pred, target)
+
+        #calculate accuracy
+        accuracy = self.accuracy(pred, target)
+
+        #log metrics
+        self.log_dict(
+            {"train_loss": loss, "train_acc": accuracy},
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+        )
+        return loss
+    
+
+    def validation_step(self, batch, batch_idx):
+        data, target = batch
+
+        #forward pass
+        pred = self.forward(data)
+
+        #calculate loss
+        loss = self.criterion(pred, target)
+
+        #calculate accuracy
+        accuracy = self.accuracy(pred, target)
+
+        #log metrics
+        self.log_dict(
+            {"val_loss": loss, "val_acc": accuracy},
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+        )
+        return loss
+
+    def test_step(self, batch, batch_idx):
+        data, target = batch
+
+        #forward pass
+        pred = self.forward(data)
+        argmax_pred = pred.argmax(dim=1).cpu() # why cpu here when down
+
+        #calculate loss
+        loss = self.criterion(pred, target)
+
+        #calculate accuracy
+        accuracy = self.accuracy(pred, target)
+
+        #update confusion matrix
+        self.confusion_matrix.update(pred, target)
+
+        #log metrics
+        self.log_dict(
+            {"test_loss": loss, "test_acc": accuracy},
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+        )
+
+        #store the predictions. labels and incorrect predictions
+
+        #converting to cpu
+        data, target, pred, argmax_pred = data.cpu(), target.cpu(), pred.cpu(), argmax_pred.cpu()
+
+        #storing the predictions
+        self.pred_store["test_preds"] = torch.cat((self.pred_store["test_preds"], argmax_pred), dim=0)
+        self.pred_store["test_labels"] = torch.cat((self.pred_store["test_labels"], target), dim=0)
+
+        for d, t, p, o in zip(data, target, argmax_pred, pred):
+            if p.eq(t.view_as(p)).item() == False:
+                self.pred_store["test_incorrect"].append(
+                    (d.cpu(), t, p, o[p.item()].cpu())
+                )
+
+        return loss
+    
+    def find_bestLR_LRFinder(self, optimizer):
+
+        lr_finder = LRFinder(self, optimizer, criterian = self.criterion)
+        lr_finder.range_test(
+            self.trainer.datamodule.train_dataloader(),
+            end_lr=config.LRFINDER_END_LR,
+            num_iter=config.LRFINDER_NUM_ITERATIONS,
+            step_mode=config.LRFINDER_STEP_MODE 
+        )
+        # best_lr = None
+        # Extract the loss and learning rate from history
+        loss = np.array(lr_finder.history['loss'])
+        lr = np.array(lr_finder.history['lr'])
+    
+        # Find the learning rate with steepest negative gradient
+        gradient = np.gradient(loss)
+        idx = np.argmin(gradient)
+        best_lr = lr[idx]
+        
+        try:
+            _, y = lr_finder.plot() 
+        except Exception as e:
+            pass
+
+        print("BEST_LR: ", best_lr)
+        lr_finder.reset()
+        
+        return best_lr
+    
+    def configure_optimizers(self):
+        optimizer = self.get_only_optimizer()
+        best_lr = self.find_bestLR_LRFinder(optimizer)
+        scheduler = OneCycleLR(
+            optimizer,
+            max_lr=best_lr, #used best_lr insted of hard coded values
+            steps_per_epoch=len(self.trainer.datamodule.train_dataloader()),
+            epochs=config.NUM_EPOCHS,
+            pct_start=5 / config.NUM_EPOCHS,
+            div_factor=config.OCLR_DIV_FACTOR,
+            three_phase=config.OCLR_THREE_PHASE,
+            final_div_factor=config.OCLR_FINAL_DIV_FACTOR,
+            anneal_strategy=config.OCLR_ANNEAL_STRATEGY
+        )
+
+        return [optimizer], [{"scheduler": scheduler, "interval": "step", "frequency": 1}]
+    
+    def get_only_optimizer(self):
+        optimizer = optim.Adam(
+            self.parameters(),lr=self.learning_rate, weight_decay=self.weight_decay
+        )
+        return optimizer
+    
+    def on_test_end(self) -> None:
+        super().on_test_end()
+
+        #Confusion Matrix
+        confmat = self.confusion_matrix.cpu().compute().numpy()
+        if config.NORM_CONF_MAT:
+            df_confmat = pd.DataFrame(
+                confmat / np.sum(confmat, axis=1)[:, None],
+                index=[i for i in config.CLASSES],
+                columns=[i for i in config.CLASSES],
+            )
+        else:
+            df_confmat = pd.DataFrame(
+                confmat,
+                index=[i for i in config.CLASSES],
+                columns=[i for i in config.CLASSES],
+            )
+        plt.figure(figsize=(7, 5))
+        sns.heatmap(df_confmat, annot=True, cmap="Blues", fmt=".3f", linewidths=0.5)
+        plt.tight_layout()
+        plt.ylabel("True label")
+        plt.xlabel("Predicted label")
+        plt.show()
+
+    def plot_incorrect_predictions_helper(self, num_imgs=10):
+        plot_incorrect_preds(
+            self.pred_store["test_incorrect"], config.CLASSES, num_imgs
+        )

Utilities/transforms.py

@@ -0,0 +1,11 @@
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+# Define the transforms (only test)
+
+test_transforms = A.Compose([
+
+    A.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261)),
+    ToTensorV2()
+
+])

Utilities/utils.py

@@ -0,0 +1,77 @@
+import torch
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+
+from . import config
+from .transforms import test_transforms
+
+
+def generate_confidences(
+    model,
+    input_img,
+    num_top_preds,
+):
+    input_img = test_transforms(image=input_img)
+    input_img = input_img["image"]
+
+    input_img = input_img.unsqueeze(0)
+    model.eval()
+    log_probs = model(input_img)[0].detach()
+    model.train()
+    probs = torch.exp(log_probs)
+
+    confidences = {
+        config.CLASSES[i]: float(probs[i]) for i in range(len(config.CLASSES))
+    }
+    # Select top 5 confidences based on value
+    confidences = {
+        k: v
+        for k, v in sorted(confidences.items(), key=lambda item: item[1], reverse=True)[
+            :num_top_preds
+        ]
+    }
+    return input_img, confidences
+
+
+def generate_gradcam(
+    model,
+    org_img,
+    input_img,
+    show_gradcam,
+    gradcam_layer,
+    gradcam_opacity,
+):
+    if show_gradcam:
+        if gradcam_layer == -1:
+            target_layers = [model.l3[-1]]
+        elif gradcam_layer == -2:
+            target_layers = [model.l2[-1]]
+
+        cam = GradCAM(
+            model=model,
+            target_layers=target_layers,
+        )
+        grayscale_cam = cam(input_tensor=input_img, targets=None)
+        grayscale_cam = grayscale_cam[0, :]
+
+        visualization = show_cam_on_image(
+            org_img / 255,
+            grayscale_cam,
+            use_rgb=True,
+            image_weight=(1 - gradcam_opacity),
+        )
+    else:
+        visualization = None
+    return visualization
+
+
+def generate_missclassified_imgs(
+    model,
+    show_misclassified,
+    num_misclassified,
+):
+    if show_misclassified:
+        plot = model.plot_incorrect_predictions_helper(num_misclassified)
+    else:
+        plot = None
+    return plot

Utilities/visualize.py

@@ -0,0 +1,31 @@
+import matplotlib.pyplot as plt
+from torchvision import transforms
+import random as rand
+
+def plot_incorrect_preds(incorrect, classes, num_imgs):
+    # num_imgs is a multiple of 5
+    assert num_imgs % 5 == 0
+    assert len(incorrect) >= num_imgs
+
+    incorrect_inds = rand.sample(range(len(incorrect)), num_imgs)
+
+    # incorrect (data, target, pred, output)
+    fig = plt.figure(figsize=(10, num_imgs // 2))
+    plt.suptitle("Target | Predicted Label")
+    for i in range(num_imgs):
+        cur_incorrect = incorrect[incorrect_inds[i]]
+        plt.subplot(num_imgs // 5, 5, i + 1, aspect="auto")
+
+        # unnormalize = T.Normalize((-mean / std).tolist(), (1.0 / std).tolist())
+        unnormalized = transforms.Normalize(
+            (-1.98947368, -1.98436214, -1.71072797), (4.048583, 4.11522634, 3.83141762)
+        )(cur_incorrect[i][0])
+        plt.imshow(transforms.ToPILImage()(unnormalized))
+        plt.title(
+            f"{classes[cur_incorrect[i][1].item()]}|{classes[cur_incorrect[i][2].item()]}",
+            # fontsize=8,
+        )
+        plt.xticks([])
+        plt.yticks([])
+    plt.tight_layout()
+    return fig

app.py

@@ -0,0 +1,108 @@
+import gradio as gr
+import torch
+
+from Utilities.model import Net
+from Utilities import config
+from Utilities.utils import generate_confidences, generate_gradcam, generate_missclassified_imgs
+
+inputs = [
+
+    gr.Image(shape=(32, 32), label="Input Image"),
+    gr.Slider(minimum=1, maximum=10, step=1, label="Number of Top Prediction to Display"),
+    gr.Checkbox(default=False, label="Show GradCAM"),
+    gr.Slider(minimum=-2, maximum=-1, step=1, value=-1, label="GradCAM Layer (from the end)"),
+    gr.Slider(minimum=0, maximum=1, value=0.5, label="GradCAM Heatmap Opacity"),
+    gr.Checkbox(label="Show Incorrect Predictions"),
+    gr.Slider(minimum=5, maximum=50, step=5, label="Number of Incorrect Predictions to Display"),
+
+]
+
+model = Net(
+    num_classes=config.NUM_CLASSES,
+    dropout_percentage = config.DROPOUT_PERCENTAGE,
+    norm = config.LAYER_NORM,
+    criterion = config.CRITERION,
+    learning_rate = config.LEARNING_RATE,
+    weight_decay = config.WEIGHT_DECAY
+)
+
+model.load_state_dict(
+    torch.load(
+        config.MODEL_PATH,
+        map_location=torch.device(config.ACCELERATOR)
+    )
+)
+
+model.pred_store = torch.load(config.PRED_STORE_PATH, map_location=torch.device(config.ACCELERATOR))
+
+def generate_gradio_output(
+        input_img,
+        num_top_preds,
+        show_gradcam,
+        gradcam_layer,
+        gradcam_opacity,
+        show_misclassified,
+        num_misclassified,
+):
+    processed_img, confidences = generate_confidences(
+        model=model,
+        input_img=input_img,
+        num_top_preds=num_top_preds 
+    )
+
+    visulization = generate_gradcam(
+        model=model,
+        org_img=input_img,
+        input_img=processed_img,
+        show_gradcam=show_gradcam,
+        gradcam_layer=gradcam_layer,
+        gradcam_opacity=gradcam_opacity,
+    )
+
+    plot = generate_missclassified_imgs(
+        model=model,
+        show_misclassified=show_misclassified,
+        num_misclassified=num_misclassified,
+    )
+
+    return confidences, visulization, plot
+
+outputs = [
+    gr.Label(visible=True, scale=0.5, label="Classification Confidences"),
+    gr.Image(shape=(32, 32), label="GradCAM Visualization").style(
+        width=256, height=256, visible=True
+    ),
+    gr.Plot(visible=True, label="Misclassified Images")
+]
+
+examples = [
+    [config.EXAMPLE_IMG_PATH + "cat.jpeg", 3, True, -2, 0.68, True, 40],
+    [config.EXAMPLE_IMG_PATH + "horse.jpg", 3, True, -2, 0.59, True, 25],
+    [config.EXAMPLE_IMG_PATH + "bird.webp", 10, True, -1, 0.55, True, 20],
+    [config.EXAMPLE_IMG_PATH + "dog1.jpg", 10, True, -1, 0.33, True, 45],
+    [config.EXAMPLE_IMG_PATH + "frog1.webp", 5, True, -1, 0.64, True, 40],
+    [config.EXAMPLE_IMG_PATH + "deer.webp", 1, True, -2, 0.45, True, 20],
+    [config.EXAMPLE_IMG_PATH + "airplane.png", 3, True, -2, 0.43, True, 40],
+    [config.EXAMPLE_IMG_PATH + "shipp.jpg", 7, True, -1, 0.6, True, 30],
+    [config.EXAMPLE_IMG_PATH + "car.jpg", 2, True, -1, 0.68, True, 30],
+    [config.EXAMPLE_IMG_PATH + "truck1.jpg", 5, True, -2, 0.51, True, 35],
+]
+
+title = "Image Classification (CIFAR10 - 10 Classes) with GradCAM"
+description = """A simple Gradio interface to visualize the output of a CNN trained on CIFAR10 dataset with GradCAM and Misclassified images. 
+The architecture is inspired from David Page's (myrtle.ai) DAWNBench winning model archiecture.
+Please input the image and select the number of top predictions to display - you will see the top predictions and their corresponding confidence scores.
+You can also select whether to show GradCAM for the particular image (utilizes the gradients of the classification score with respect to the final convolutional feature map, to identify the parts of an input image that most impact the classification score).
+You need to select the model layer where the gradients need to be plugged from - this affects how much of the image is used to compute the GradCAM.
+You can also select whether to show misclassified images - these are the images that the model misclassified.
+Some examples are provided in the examples tab.
+"""
+
+gr.Interface(
+    fn=generate_gradio_output,
+    inputs=inputs,
+    outputs=outputs,
+    title=title,
+    description=description,
+    examples=examples
+).launch()