Upload resnet.py

resnet.py

@@ -0,0 +1,520 @@
+#!/usr/bin/env python3
+"""
+PyTorch Lightning for ResNet Architecture
+Author: Shilpaj Bhalerao
+"""
+# Standard Library Imports
+import os
+import math
+
+# Third-Party Imports
+import numpy as np
+import matplotlib.pyplot as plt
+import albumentations as A
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import DataLoader, random_split
+
+from torchvision import transforms
+from torchvision.datasets import CIFAR10
+
+from pytorch_lightning import LightningModule, Trainer
+from torchmetrics import Accuracy
+
+# Local Imports
+from datasets import AlbumDataset
+from utils import get_cifar_statistics
+from visualize import visualize_cifar_augmentation, display_cifar_data_samples
+
+
+class Layers:
+    """
+    Class containing different types of Convolutional layer
+    """
+
+    def __init__(self, groups=1):
+        """
+        Constructor
+        """
+        self.group = groups
+
+    @staticmethod
+    def standard_conv_layer(in_channels: int,
+                            out_channels: int,
+                            kernel_size: int = 3,
+                            padding: int = 0,
+                            stride: int = 1,
+                            dilation: int = 1,
+                            normalization: str = "batch",
+                            last_layer: bool = False,
+                            conv_type: str = "standard",
+                            groups: int = 1):
+        """
+        Method to return a standard convolution block
+        :param in_channels: Number of input channels
+        :param out_channels: Number of output channels
+        :param kernel_size: Size of the kernel used in the layer
+        :param padding: Padding used in the layer
+        :param stride: Stride used for convolution
+        :param dilation: Dilation for Atrous convolution
+        :param normalization: Type of normalization technique used
+        :param last_layer: Flag to indicate if the layer is last convolutional layer of the network
+        :param conv_type: Type of convolutional layer
+        :param groups: Number of Groups for Group Normalization
+        """
+        # Select normalization type
+        if normalization == "layer":
+            _norm_layer = nn.GroupNorm(1, out_channels)
+        elif normalization == "group":
+            if not groups:
+                raise ValueError("Value of group is not defined")
+            _norm_layer = nn.GroupNorm(groups, out_channels)
+        else:
+            _norm_layer = nn.BatchNorm2d(out_channels)
+
+        # Select the convolution layer type
+        if conv_type == "standard":
+            conv_layer = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, stride=stride,
+                                   kernel_size=kernel_size, bias=False, padding=padding)
+        elif conv_type == "depthwise":
+            conv_layer = Layers.depthwise_conv(in_channels=in_channels, out_channels=out_channels, stride=stride,
+                                               padding=padding)
+        elif conv_type == "dilated":
+            conv_layer = Layers.dilated_conv(in_channels=in_channels, out_channels=out_channels, stride=stride,
+                                             padding=padding, dilation=dilation)
+
+        # For last layer only return the convolution output
+        if last_layer:
+            return nn.Sequential(conv_layer)
+        return nn.Sequential(
+            conv_layer,
+            _norm_layer,
+            nn.ReLU(),
+            # nn.Dropout(self.dropout_value)
+        )
+
+    @staticmethod
+    def resnet_block(channels):
+        """
+        Method to create a RESNET block
+        """
+        return nn.Sequential(
+            nn.Conv2d(in_channels=channels, out_channels=channels, stride=1, kernel_size=3, bias=False, padding=1),
+            nn.BatchNorm2d(channels),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=channels, out_channels=channels, stride=1, kernel_size=3, bias=False, padding=1),
+            nn.BatchNorm2d(channels),
+            nn.ReLU(),
+        )
+
+    @staticmethod
+    def custom_block(input_channels, output_channels):
+        """
+        Method to create a custom configured block
+        :param input_channels: Number of input channels
+        :param output_channels: Number of output channels
+        """
+        return nn.Sequential(
+            nn.Conv2d(in_channels=input_channels, out_channels=output_channels, stride=1, kernel_size=3, bias=False,
+                      padding=1),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.BatchNorm2d(output_channels),
+            nn.ReLU(),
+        )
+
+    @staticmethod
+    def depthwise_conv(in_channels, out_channels, stride=1, padding=0):
+        """
+        Method to return the depthwise separable convolution layer
+        :param in_channels: Number of input channels
+        :param out_channels: Number of output channels
+        :param padding: Padding used in the layer
+        :param stride: Stride used for convolution
+        """
+        return nn.Sequential(
+            nn.Conv2d(in_channels=in_channels, out_channels=in_channels, stride=stride, groups=in_channels,
+                      kernel_size=3, bias=False, padding=padding),
+            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, stride=stride, kernel_size=1, bias=False,
+                      padding=0)
+        )
+
+    @staticmethod
+    def dilated_conv(in_channels, out_channels, stride=1, padding=0, dilation=1):
+        """
+        Method to return the dilated convolution layer
+        :param in_channels: Number of input channels
+        :param out_channels: Number of output channels
+        :param stride: Stride used for convolution
+        :param padding: Padding used in the layer
+        :param dilation: Dilation value for a kernel
+        """
+        return nn.Sequential(
+            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, stride=stride, kernel_size=3, bias=False,
+                      padding=padding, dilation=dilation)
+        )
+
+
+class LITResNet(LightningModule, Layers):
+    """
+    David's Model Architecture for Session-10 CIFAR10 dataset
+    """
+
+    def __init__(self, class_names, data_dir='/data/'):
+        """
+        Constructor
+        """
+        # Initialize the Module class
+        super().__init__()
+
+        # Initialize variables
+        self.classes = class_names
+        self.data_dir = data_dir
+        self.num_classes = 10
+        self._learning_rate = 0.03
+        self.inv_normalize = 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]
+        )
+        self.batch_size = 512
+        self.epochs = 24
+        self.accuracy = Accuracy(task='multiclass',
+                                 num_classes=10)
+        self.train_transforms = transforms.Compose([transforms.ToTensor()])
+        self.test_transforms = transforms.Compose([transforms.ToTensor()])
+        self.stats_train = None
+        self.stats_test = None
+        self.cifar10_train = None
+        self.cifar10_test = None
+        self.cifar10_val = None
+        self.misclassified_data = None
+
+        # Defined Layers for the model
+        self.prep_layer = None
+        self.custom_block1 = None
+        self.custom_block2 = None
+        self.custom_block3 = None
+        self.resnet_block1 = None
+        self.resnet_block3 = None
+        self.pool4 = None
+        self.fc = None
+        self.dropout_value = None
+
+        # Initialize all the layers
+        self.model_layers()
+
+    # ##################################################################################################
+    # ################################ Model Architecture Related Hooks ################################
+    # ##################################################################################################
+    def model_layers(self):
+        """
+        Method to initialize layers for the model
+        """
+        # Prep Layer
+        self.prep_layer = Layers.standard_conv_layer(in_channels=3, out_channels=64, kernel_size=3, padding=1, stride=1)
+
+        # Convolutional Block-1
+        self.custom_block1 = Layers.custom_block(input_channels=64, output_channels=128)
+        self.resnet_block1 = Layers.resnet_block(channels=128)
+
+        # Convolutional Block-2
+        self.custom_block2 = Layers.custom_block(input_channels=128, output_channels=256)
+
+        # Convolutional Block-3
+        self.custom_block3 = Layers.custom_block(input_channels=256, output_channels=512)
+        self.resnet_block3 = Layers.resnet_block(channels=512)
+
+        # MaxPool Layer
+        self.pool4 = nn.MaxPool2d(kernel_size=4, stride=2)
+
+        # Fully Connected Layer
+        self.fc = nn.Linear(in_features=512, out_features=10, bias=False)
+
+        # Dropout value of 10%
+        self.dropout_value = 0.1
+
+    def forward(self, x):
+        """
+        Forward pass for model training
+        :param x: Input layer
+        :return: Model Prediction
+        """
+        # Prep Layer
+        x = self.prep_layer(x)
+
+        # Convolutional Block-1
+        x = self.custom_block1(x)
+        r1 = self.resnet_block1(x)
+        x = x + r1
+
+        # Convolutional Block-2
+        x = self.custom_block2(x)
+
+        # Convolutional Block-3
+        x = self.custom_block3(x)
+        r2 = self.resnet_block3(x)
+        x = x + r2
+
+        # MaxPool Layer
+        x = self.pool4(x)
+
+        # Fully Connected Layer
+        x = x.view(-1, 512)
+        x = self.fc(x)
+
+        return F.log_softmax(x, dim=1)
+
+    # ##################################################################################################
+    # ############################## Training Configuration Related Hooks ##############################
+    # ##################################################################################################
+
+    def configure_optimizers(self):
+        """
+        Method to configure the optimizer and learning rate scheduler
+        """
+        learning_rate = 0.03
+        weight_decay = 1e-4
+        optimizer = optim.Adam(self.parameters(), lr=learning_rate, weight_decay=weight_decay)
+
+        # Scheduler
+        scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer,
+                                                        max_lr=self._learning_rate,
+                                                        steps_per_epoch=len(self.train_dataloader()),
+                                                        epochs=self.epochs,
+                                                        pct_start=5 / self.epochs,
+                                                        div_factor=100,
+                                                        three_phase=False,
+                                                        final_div_factor=100,
+                                                        anneal_strategy="linear"
+                                                        )
+        return [optimizer], [{'scheduler': scheduler, 'interval': 'step'}]
+
+    @property
+    def learning_rate(self) -> float:
+        """
+        Method to get the learning rate value
+        """
+        return self._learning_rate
+
+    @learning_rate.setter
+    def learning_rate(self, value: float):
+        """
+        Method to set the learning rate value
+        :param value: Updated value of learning rate
+        """
+        self._learning_rate = value
+
+    def set_training_confi(self, *, epochs, batch_size):
+        """
+        Method to set parameters required for model training
+        :param epochs: Number of epochs for which model is to be trained
+        :param batch_size: Batch Size
+        """
+        self.epochs = epochs
+        self.batch_size = batch_size
+
+    # #################################################################################################
+    # ################################## Training Loop Related Hooks ##################################
+    # #################################################################################################
+    def training_step(self, train_batch, batch_index):
+        """
+        Method called on training dataset to train the model
+        :param train_batch: Batch containing images and labels
+        :param batch_index: Index of the batch
+        """
+        x, y = train_batch
+        logits = self.forward(x)
+        loss = F.cross_entropy(logits, y)
+        preds = torch.argmax(logits, dim=1)
+        self.accuracy(preds, y)
+
+        self.log("train_loss", loss, prog_bar=True)
+        self.log("train_acc", self.accuracy, prog_bar=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        """
+        Method called on validation dataset to check if the model is learning
+        :param batch: Batch containing images and labels
+        :param batch_idx: Index of the batch
+        """
+        x, y = batch
+        logits = self.forward(x)
+        loss = F.nll_loss(logits, y)
+        preds = torch.argmax(logits, dim=1)
+        self.accuracy(preds, y)
+
+        # Calling self.log will surface up scalars for you in TensorBoard
+        self.log("val_loss", loss, prog_bar=True)
+        self.log("val_acc", self.accuracy, prog_bar=True)
+        return loss
+
+    def test_step(self, batch, batch_idx):
+        """
+        Method called on test dataset to check model performance on unseen data
+        :param batch: Batch containing images and labels
+        :param batch_idx: Index of the batch
+        """
+        # Here we just reuse the validation_step for testing
+        return self.validation_step(batch, batch_idx)
+
+    # ##############################################################################################
+    # ##################################### Data Related Hooks #####################################
+    # ##############################################################################################
+
+    def set_transforms(self, train_set_transforms: dict, test_set_transforms: dict):
+        """
+        Method to set the transformations to be done on training and test datasets
+        :param train_set_transforms: Dictionary of transformations for training dataset
+        :param test_set_transforms: Dictionary of transformations for test dataset
+        """
+        self.train_transforms = A.Compose(train_set_transforms.values())
+        self.test_transforms = A.Compose(test_set_transforms.values())
+
+    def prepare_data(self):
+        """
+        Method to download the dataset
+        """
+        self.stats_train = CIFAR10('./data', train=True, download=True, transform=transforms.ToTensor())
+        self.stats_test = CIFAR10('./data', train=False, download=True, transform=transforms.ToTensor())
+
+    def setup(self, stage=None):
+        """
+        Method to create Split the dataset into train, test and val
+        """
+        # Only if dataset is not already split, perform the split operation
+        if not self.cifar10_train and not self.cifar10_test and not self.cifar10_val:
+
+            # Assign train/val datasets for use in dataloaders
+            if stage == "fit" or stage is None:
+                cifar10_full = AlbumDataset(self.data_dir, train=True, download=True, transform=self.train_transforms)
+                self.cifar10_train, self.cifar10_val = random_split(cifar10_full, [45_000, 5_000])
+
+            # Assign test dataset for use in dataloader(s)
+            if stage == "test" or stage is None:
+                self.cifar10_test = AlbumDataset(self.data_dir, train=False, download=True,
+                                                 transform=self.test_transforms)
+
+    def train_dataloader(self):
+        """
+        Method to return the DataLoader for Training set
+        """
+        return DataLoader(self.cifar10_train, batch_size=self.batch_size, num_workers=os.cpu_count())
+
+    def val_dataloader(self):
+        """
+        Method to return the DataLoader for the Validation set
+        """
+        return DataLoader(self.cifar10_val, batch_size=self.batch_size, num_workers=os.cpu_count())
+
+    def test_dataloader(self):
+        """
+        Method to return the DataLoader for the Test set
+        """
+        return DataLoader(self.cifar10_test, batch_size=self.batch_size, num_workers=os.cpu_count())
+
+    def get_statistics(self, data_set_type="Train"):
+        """
+        Method to get the statistics for CIFAR10 dataset
+        """
+        # Execute self.prepare_data() only if not done earlier
+        if not self.stats_train and not self.stats_test:
+            self.prepare_data()
+
+        # Print stats for selected dataset
+        if data_set_type == "Train":
+            get_cifar_statistics(self.stats_train)
+        else:
+            get_cifar_statistics(self.stats_test, data_set_type="Test")
+
+    def display_data_samples(self, dataset="train", num_of_images=20):
+        """
+        Method to display data samples
+        """
+        # Execute self.prepare_data() only if not done earlier
+        try:
+            assert self.stats_train
+        except AttributeError:
+            self.prepare_data()
+
+        if dataset == "train":
+            display_cifar_data_samples(self.stats_train, num_of_images, self.classes)
+        else:
+            display_cifar_data_samples(self.stats_test, num_of_images, self.classes)
+
+    @staticmethod
+    def visualize_augmentation(aug_set_transforms: dict):
+        """
+        Method to visualize augmentations
+        :param aug_set_transforms: Dictionary of transformations to be visualized
+        """
+        aug_train = AlbumDataset('./data', train=True, download=True)
+        visualize_cifar_augmentation(aug_train, aug_set_transforms)
+
+    # #############################################################################################
+    # ############################## Misclassified Data Related Hooks ##############################
+    # #############################################################################################
+
+    def get_misclassified_data(self):
+        """
+        Function to run the model on test set and return misclassified images
+        """
+        if self.misclassified_data:
+            return self.misclassified_data
+
+        self.misclassified_data = []
+        self.prepare_data()
+        self.setup()
+
+        test_loader = self.test_dataloader()
+
+        # 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(self.device), target.to(self.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 = self.forward(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:
+                        self.misclassified_data.append((image, label, pred))
+        return self.misclassified_data
+
+    def display_cifar_misclassified_data(self, number_of_samples: int = 10):
+        """
+        Function to plot images with labels
+        :param number_of_samples: Number of images to print
+        """
+        if not self.misclassified_data:
+            self.misclassified_data = self.get_misclassified_data()
+
+        fig = plt.figure(figsize=(10, 10))
+
+        x_count = 5
+        y_count = 1 if number_of_samples <= 5 else math.floor(number_of_samples / x_count)
+
+        for i in range(number_of_samples):
+            plt.subplot(y_count, x_count, i + 1)
+            img = self.misclassified_data[i][0].squeeze().to('cpu')
+            img = self.inv_normalize(img)
+            plt.imshow(np.transpose(img, (1, 2, 0)))
+            plt.title(
+                r"Correct: " + self.classes[self.misclassified_data[i][1].item()] + '\n' + 'Output: ' + self.classes[
+                    self.misclassified_data[i][2].item()])
+            plt.xticks([])
+            plt.yticks([])