Upload 8 files

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)

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(),
+    ]
+)