Upload 8 files

datasets/__init__.py

@@ -0,0 +1 @@
+from .cifar10 import CIFAR10

datasets/cifar10.py

@@ -0,0 +1,32 @@
+import numpy as np
+import cv2
+from torchvision import datasets
+import albumentations as A
+
+from .generic import MyDataSet
+
+
+class AlbCIFAR10(datasets.CIFAR10):
+    def __init__(self, root, alb_transform=None, **kwargs):
+        super(AlbCIFAR10, self).__init__(root, **kwargs)
+        self.alb_transform = alb_transform
+
+    def __getitem__(self, index):
+        image, label = super(AlbCIFAR10, self).__getitem__(index)
+        if self.alb_transform is not None:
+            image = self.alb_transform(image=np.array(image))['image']
+        return image, label
+
+
+class CIFAR10(MyDataSet):
+    DataSet = AlbCIFAR10
+    mean = (0.49139968, 0.48215827, 0.44653124)
+    std = (0.24703233, 0.24348505, 0.26158768)
+    default_alb_transforms = [
+        A.ToGray(p=0.2),
+        A.PadIfNeeded(40, 40, p=1),
+        A.RandomCrop(32, 32, p=1),
+        A.HorizontalFlip(p=0.5),
+        # Since normalisation was the first step, mean is already 0, so cutout fill_value = 0
+        A.CoarseDropout(max_holes=1, max_height=8, max_width=8, fill_value=0, p=1)
+    ]

datasets/generic.py

@@ -0,0 +1,111 @@
+import os
+from abc import ABC
+from functools import cached_property
+
+import torch
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+try:
+    from epoch.utils import plot_examples
+except ModuleNotFoundError:
+    from utils import plot_examples
+
+
+class MyDataSet(ABC):
+    DataSet = None
+    mean = None
+    std = None
+    classes = None
+    default_alb_transforms = None
+
+    def __init__(self, batch_size=1, normalize=True, shuffle=True, augment=True, alb_transforms=None):
+        self.batch_size = batch_size
+        self.normalize = normalize
+        self.shuffle = shuffle
+        self.augment = augment
+        self.alb_transforms = alb_transforms or self.default_alb_transforms
+
+        self.loader_kwargs = {'batch_size': batch_size, 'num_workers': os.cpu_count(), 'pin_memory': True}
+
+    @classmethod
+    def set_classes(cls, data):
+        if cls.classes is None:
+            cls.classes = {i: c for i, c in enumerate(data.classes)}
+
+    @cached_property
+    def train_data(self):
+        res = self.DataSet('../data', train=True, download=True, alb_transform=self.get_train_transforms())
+        self.set_classes(res)
+        return res
+
+    @cached_property
+    def test_data(self):
+        res = self.DataSet('../data', train=False, download=True, alb_transform=self.get_test_transforms())
+        self.set_classes(res)
+        return res
+
+    @cached_property
+    def train_loader(self):
+        return torch.utils.data.DataLoader(self.train_data, shuffle=self.shuffle, **self.loader_kwargs)
+
+    @cached_property
+    def test_loader(self):
+        return torch.utils.data.DataLoader(self.test_data, shuffle=False, **self.loader_kwargs)
+
+    @cached_property
+    def example_iter(self):
+        return iter(self.train_loader)
+
+    def get_train_transforms(self):
+        all_transforms = list()
+        if self.normalize:
+            all_transforms.append(A.Normalize(self.mean, self.std))
+        if self.augment and self.alb_transforms is not None:
+            all_transforms.extend(self.alb_transforms)
+        all_transforms.append(ToTensorV2())
+        return A.Compose(all_transforms)
+
+    def get_test_transforms(self):
+        all_transforms = list()
+        if self.normalize:
+            all_transforms.append(A.Normalize(self.mean, self.std))
+        all_transforms.append(ToTensorV2())
+        return A.Compose(all_transforms)
+
+    def download(self):
+        self.DataSet('../data', train=True, download=True)
+        self.DataSet('../data', train=False, download=True)
+
+    def denormalise(self, tensor):
+        result = tensor.clone().detach().requires_grad_(False)
+        if self.normalize:
+            for t, m, s in zip(result, self.mean, self.std):
+                t.mul_(s).add_(m)
+        return result
+
+    def show_transform(self, img):
+        if self.normalize:
+            img = self.denormalise(img)
+        if len(self.mean) == 3:
+            return img.permute(1, 2, 0)
+        else:
+            return img.squeeze(0)
+
+    def show_examples(self, figsize=(8, 6)):
+        batch_data, batch_label = next(self.example_iter)
+        images = list()
+        labels = list()
+
+        for i in range(len(batch_data)):
+            image = batch_data[i]
+            image = self.show_transform(image)
+
+            label = batch_label[i].item()
+            if self.classes is not None:
+                label = f'{label}:{self.classes[label]}'
+
+            images.append(image)
+            labels.append(label)
+
+        plot_examples(images, labels, figsize=figsize)

models/custom_resnet.py

@@ -0,0 +1,159 @@
+from torch import nn
+from torch import optim
+from pytorch_lightning import LightningModule
+from torchmetrics import MeanMetric
+from torch_lr_finder import LRFinder
+
+from utils.metrics import RunningAccuracy
+
+
+class ConvLayer(nn.Module):
+    def __init__(self, input_c, output_c, bias=False, stride=1, padding=1, pool=False, dropout=0.):
+        super(ConvLayer, self).__init__()
+
+        layers = list()
+        layers.append(
+            nn.Conv2d(input_c, output_c, kernel_size=3, bias=bias, stride=stride, padding=padding,
+                      padding_mode='replicate')
+        )
+        if pool:
+            layers.append(nn.MaxPool2d(kernel_size=2, stride=2))
+        layers.append(nn.BatchNorm2d(output_c))
+        layers.append(nn.ReLU())
+        if dropout > 0:
+            layers.append(nn.Dropout(dropout))
+
+        self.all_layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.all_layers(x)
+
+
+class CustomLayer(nn.Module):
+    def __init__(self, input_c, output_c, pool=True, residue=2, dropout=0.):
+        super(CustomLayer, self).__init__()
+
+        self.pool_block = ConvLayer(input_c, output_c, pool=pool, dropout=dropout)
+        self.res_block = None
+        if residue > 0:
+            layers = list()
+            for i in range(0, residue):
+                layers.append(ConvLayer(output_c, output_c, pool=False, dropout=dropout))
+            self.res_block = nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.pool_block(x)
+        if self.res_block is not None:
+            x_ = x
+            x = self.res_block(x)
+            # += operator causes inplace errors in pytorch if done right after relu.
+            x = x + x_
+        return x
+
+
+class Model(LightningModule):
+    def __init__(self, dataset, dropout=0.05, max_epochs=24):
+        super(Model, self).__init__()
+
+        self.dataset = dataset
+
+        self.network = nn.Sequential(
+            CustomLayer(3, 64, pool=False, residue=0, dropout=dropout),
+            CustomLayer(64, 128, pool=True, residue=2, dropout=dropout),
+            CustomLayer(128, 256, pool=True, residue=0, dropout=dropout),
+            CustomLayer(256, 512, pool=True, residue=2, dropout=dropout),
+            nn.MaxPool2d(kernel_size=4, stride=4),
+            nn.Flatten(),
+            nn.Linear(512, 10)
+        )
+
+        self.criterion = nn.CrossEntropyLoss()
+        self.train_accuracy = RunningAccuracy()
+        self.val_accuracy = RunningAccuracy()
+        self.train_loss = MeanMetric()
+        self.val_loss = MeanMetric()
+
+        self.max_epochs = max_epochs
+        self.epoch_counter = 1
+
+    def forward(self, x):
+        return self.network(x)
+
+    def common_step(self, batch, loss_metric, acc_metric):
+        x, y = batch
+        batch_len = y.numel()
+        logits = self.forward(x)
+        loss = self.criterion(logits, y)
+        loss_metric.update(loss, batch_len)
+        acc_metric.update(logits, y)
+        return loss
+
+    def training_step(self, batch, batch_idx):
+        return self.common_step(batch, self.train_loss, self.train_accuracy)
+
+    def on_train_epoch_end(self):
+        print(f"Epoch: {self.epoch_counter}, Train: Loss: {self.train_loss.compute():0.4f}, Accuracy: "
+              f"{self.train_accuracy.compute():0.2f}")
+        self.train_loss.reset()
+        self.train_accuracy.reset()
+        self.epoch_counter += 1
+
+    def validation_step(self, batch, batch_idx):
+        loss = self.common_step(batch, self.val_loss, self.val_accuracy)
+        self.log("val_step_loss", self.val_loss, prog_bar=True, logger=True)
+        self.log("val_step_acc", self.val_accuracy, prog_bar=True, logger=True)
+        return loss
+
+    def on_validation_epoch_end(self):
+        print(f"Epoch: {self.epoch_counter}, Valid: Loss: {self.val_loss.compute():0.4f}, Accuracy: "
+              f"{self.val_accuracy.compute():0.2f}")
+        self.val_loss.reset()
+        self.val_accuracy.reset()
+
+    def predict_step(self, batch, batch_idx, dataloader_idx=0):
+        if isinstance(batch, list):
+            x, _ = batch
+        else:
+            x = batch
+        return self.forward(x)
+
+    def find_lr(self, optimizer):
+        lr_finder = LRFinder(self, optimizer, self.criterion)
+        lr_finder.range_test(self.dataset.train_loader, end_lr=0.1, num_iter=100, step_mode='exp')
+        _, best_lr = lr_finder.plot()
+        lr_finder.reset()
+        return best_lr
+
+    def configure_optimizers(self):
+        optimizer = optim.Adam(self.parameters(), lr=1e-7, weight_decay=1e-2)
+        best_lr = self.find_lr(optimizer)
+        scheduler = optim.lr_scheduler.OneCycleLR(
+            optimizer,
+            max_lr=best_lr,
+            steps_per_epoch=len(self.dataset.train_loader),
+            epochs=self.max_epochs,
+            pct_start=5/self.max_epochs,
+            div_factor=100,
+            three_phase=False,
+            final_div_factor=100,
+            anneal_strategy='linear'
+        )
+        return {
+            'optimizer': optimizer,
+            'lr_scheduler': {
+                "scheduler": scheduler,
+                "interval": "step",
+            }
+        }
+
+    def prepare_data(self):
+        self.dataset.download()
+
+    def train_dataloader(self):
+        return self.dataset.train_loader
+
+    def val_dataloader(self):
+        return self.dataset.test_loader
+
+    def predict_dataloader(self):
+        return self.val_dataloader()

utils/__init__.py

@@ -0,0 +1 @@
+from .misc import *

utils/metrics.py

@@ -0,0 +1,19 @@
+import torch
+from torch import Tensor
+from torchmetrics import Metric
+
+
+class RunningAccuracy(Metric):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.add_state("correct", default=torch.tensor(0), dist_reduce_fx="sum")
+        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")
+
+    def update(self, preds: Tensor, target: Tensor):
+        preds = preds.argmax(dim=1)
+        total = target.numel()
+        self.correct += preds.eq(target).sum()
+        self.total += total
+
+    def compute(self):
+        return 100 * self.correct.float() / self.total

utils/misc.py

@@ -0,0 +1,69 @@
+import torch
+import torchinfo
+from matplotlib import pyplot as plt
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
+
+SEED = 42
+DEVICE = None
+
+
+def get_device():
+    global DEVICE
+    if DEVICE is not None:
+        return DEVICE
+
+    if torch.cuda.is_available():
+        DEVICE = "cuda"
+    elif torch.backends.mps.is_available():
+        DEVICE = "mps"
+    else:
+        DEVICE = "cpu"
+    print("Device Selected:", DEVICE)
+    return DEVICE
+
+
+def set_seed(seed=SEED):
+    torch.manual_seed(seed)
+    if get_device() == 'cuda':
+        torch.cuda.manual_seed(seed)
+
+
+def plot_examples(images, labels, figsize=None, n=20):
+    _ = plt.figure(figsize=figsize)
+
+    for i in range(n):
+        plt.subplot(4, n//4, i + 1)
+        plt.tight_layout()
+        image = images[i]
+        plt.imshow(image, cmap='gray')
+        label = labels[i]
+        plt.title(str(label))
+        plt.xticks([])
+        plt.yticks([])
+
+
+def get_incorrect_preds(prediction, labels):
+    prediction = prediction.argmax(dim=1)
+    indices = prediction.ne(labels).nonzero().reshape(-1).tolist()
+    return indices, prediction[indices].tolist(), labels[indices].tolist()
+
+
+def get_cam_visualisation(model, dataset, input_tensor, label, target_layer, use_cuda=False):
+    grad_cam = GradCAM(model=model, target_layers=[target_layer], use_cuda=use_cuda)
+
+    targets = [ClassifierOutputTarget(label)]
+
+    grayscale_cam = grad_cam(input_tensor=input_tensor.unsqueeze(0), targets=targets)
+    # In this example grayscale_cam has only one image in the batch:
+    grayscale_cam = grayscale_cam[0, :]
+
+    output = show_cam_on_image(dataset.show_transform(input_tensor).cpu().numpy(), grayscale_cam,
+                               use_rgb=True)
+    return output
+
+
+def model_summary(model, input_size=None):
+    return torchinfo.summary(model, input_size=input_size, depth=5,
+                             col_names=["input_size", "output_size", "num_params", "params_percent"])