autoaugment experiment

Author: rasbt

README.md

@@ -288,6 +288,13 @@ Please note that the following notebooks below provide reference implementations
 
 
 
+## Data Augmentation
+
+| Title                      | Dataset | Description | Notebooks                                                    |
+| -------------------------- | ------- | ----------- | ------------------------------------------------------------ |
+| AutoAugment for Image Data | CIFAR-10     | Trains a ResNet-18 using AutoAugment         | [![PyTorch Lightning](https://img.shields.io/badge/PyTorch-Lightning-blueviolet)](pytorch-lightning_ipynb/data-augumentation/autoaugment/) |
+
+
 
 
 ## Tips and Tricks

pytorch-lightning_ipynb/data-augmentation/autoaugment/helper_utilities.py

@@ -0,0 +1,195 @@
+import lightning as L
+import matplotlib.pyplot as plt
+import pandas as pd
+import torch
+import torch.nn.functional as F
+import torchmetrics
+from torch.utils.data import DataLoader
+from torch.utils.data.dataset import random_split
+from torchvision import datasets, transforms
+
+
+class LightningModel(L.LightningModule):
+    def __init__(self, model, learning_rate):
+        super().__init__()
+
+        self.learning_rate = learning_rate
+        self.model = model
+
+        self.save_hyperparameters(ignore=["model"])
+
+        self.train_acc = torchmetrics.Accuracy(task="multiclass", num_classes=10)
+        self.val_acc = torchmetrics.Accuracy(task="multiclass", num_classes=10)
+        self.test_acc = torchmetrics.Accuracy(task="multiclass", num_classes=10)
+
+    def forward(self, x):
+        return self.model(x)
+
+    def _shared_step(self, batch):
+        features, true_labels = batch
+        logits = self(features)
+
+        loss = F.cross_entropy(logits, true_labels)
+        predicted_labels = torch.argmax(logits, dim=1)
+        return loss, true_labels, predicted_labels
+
+    def training_step(self, batch, batch_idx):
+        loss, true_labels, predicted_labels = self._shared_step(batch)
+
+        self.log("train_loss", loss)
+        self.train_acc(predicted_labels, true_labels)
+        self.log(
+            "train_acc", self.train_acc, prog_bar=True, on_epoch=True, on_step=False
+        )
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        loss, true_labels, predicted_labels = self._shared_step(batch)
+
+        self.log("val_loss", loss, prog_bar=True)
+        self.val_acc(predicted_labels, true_labels)
+        self.log("val_acc", self.val_acc, prog_bar=True)
+
+    def test_step(self, batch, batch_idx):
+        loss, true_labels, predicted_labels = self._shared_step(batch)
+        self.test_acc(predicted_labels, true_labels)
+        self.log("test_acc", self.test_acc)
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.SGD(self.parameters(), lr=self.learning_rate)
+        return optimizer
+
+
+class Cifar10DataModule(L.LightningDataModule):
+    def __init__(
+        self, data_path="./", batch_size=64, num_workers=0, height_width=(32, 32),
+        train_transform=None, test_transform=None
+    ):
+        super().__init__()
+        self.batch_size = batch_size
+        self.data_path = data_path
+        self.num_workers = num_workers
+        self.height_width = height_width
+        self.train_transform = train_transform
+        self.test_transform = test_transform
+
+    def prepare_data(self):
+        datasets.CIFAR10(root=self.data_path, download=True)
+
+        if self.train_transform is None:
+            self.train_transform = transforms.Compose(
+                [
+                    transforms.Resize(self.height_width),
+                    transforms.ToTensor(),
+                ]
+            )
+
+        if self.test_transform is None:
+            self.test_transform = transforms.Compose(
+                [
+                    transforms.Resize(self.height_width),
+                    transforms.ToTensor(),
+                ]
+            )
+            return
+
+    def setup(self, stage=None):
+        train = datasets.CIFAR10(
+            root=self.data_path,
+            train=True,
+            transform=self.train_transform,
+            download=False,
+        )
+
+        self.test = datasets.CIFAR10(
+            root=self.data_path,
+            train=False,
+            transform=self.test_transform,
+            download=False,
+        )
+
+        self.train, self.valid = random_split(train, lengths=[45000, 5000])
+
+    def train_dataloader(self):
+        train_loader = DataLoader(
+            dataset=self.train,
+            batch_size=self.batch_size,
+            drop_last=True,
+            shuffle=True,
+            num_workers=self.num_workers,
+        )
+        return train_loader
+
+    def val_dataloader(self):
+        valid_loader = DataLoader(
+            dataset=self.valid,
+            batch_size=self.batch_size,
+            drop_last=False,
+            shuffle=False,
+            num_workers=self.num_workers,
+        )
+        return valid_loader
+
+    def test_dataloader(self):
+        test_loader = DataLoader(
+            dataset=self.test,
+            batch_size=self.batch_size,
+            drop_last=False,
+            shuffle=False,
+            num_workers=self.num_workers,
+        )
+        return test_loader
+
+
+def plot_val_acc(
+    log_dir, acc_ylim=(0.5, 1.0), save_loss=None, save_acc=None):
+
+    metrics = pd.read_csv(f"{log_dir}/metrics.csv")
+
+    aggreg_metrics = []
+    agg_col = "epoch"
+
+    for i, dfg in metrics.groupby(agg_col):
+        agg = dict(dfg.mean())
+        agg[agg_col] = i
+        aggreg_metrics.append(agg)
+
+    df_metrics = pd.DataFrame(aggreg_metrics)
+    df_metrics[["val_acc"]].plot(
+        grid=True, legend=True, xlabel="Epoch", ylabel="ACC"
+    )
+
+    plt.ylim(acc_ylim)
+    if save_acc is not None:
+        plt.savefig(save_acc)
+
+        
+def plot_loss_and_acc(
+    log_dir, loss_ylim=(0.0, 0.9), acc_ylim=(0.3, 1.0), save_loss=None, save_acc=None
+):
+
+    metrics = pd.read_csv(f"{log_dir}/metrics.csv")
+
+    aggreg_metrics = []
+    agg_col = "epoch"
+    for i, dfg in metrics.groupby(agg_col):
+        agg = dict(dfg.mean())
+        agg[agg_col] = i
+        aggreg_metrics.append(agg)
+
+    df_metrics = pd.DataFrame(aggreg_metrics)
+    df_metrics[["train_loss"]].plot(
+        grid=True, legend=True, xlabel="Epoch", ylabel="Loss"
+    )
+
+    plt.ylim(loss_ylim)
+    if save_loss is not None:
+        plt.savefig(save_loss)
+
+    df_metrics[["train_acc", "val_acc"]].plot(
+        grid=True, legend=True, xlabel="Epoch", ylabel="ACC"
+    )
+
+    plt.ylim(acc_ylim)
+    if save_acc is not None:
+        plt.savefig(save_acc)