Copyright (c) 2023 Massachusetts Institute of Technology

SPDX-License-Identifier: MIT

Building Workflows for Cross Validation Training and Adversarial Robustness Analysis#

This notebook demonstrates how the to build experiment workflows for configurable, repeatable, and scalable (CRS) experimentation. Two basic workflows will be demonstrated in this tutorial:

Cross-Validation Workflow: Performs cross-validation training that logs accuracy and loss across data folds
Robustness Curve Workflow: Loads a trained model and assesses the impact of adversarial perturbations on the model’s performance; the model’s performance metric is plotted against the increasing “severity” of the perturbation

Here, “workflow” has a precise meaning. In the parlance of mushin a workflow is an API for describing how we configure, launch, and post-process one or more tasks. These workflows leverage Hydra and hydra-zen so that they are highly configurable and so that each job launched by a workflow is self-documenting and reproducible. In this tutorial, we also make use of PyTorch Lightning to eliminate boilerplate code associated with training and testing a PyTorch model.

Getting Started#

We will install the rAI-toolbox and then we will create a Jupyter notebook in which we will complete this tutorial.

Installing `rai_toolbox`#

To install the toolbox (along with its mushin capabilities) in your Python environment, run the following command in your terminal:

$ pip install rai-toolbox[mushin]

To verify that the toolbox is installed as-expected, open a Python console and try importing rai_toolbox.

>>> import rai_toolbox

You will also need to install scikit-learn; please follow these instructions.

Opening a Jupyter notebook#

If you do not have Jupyter Notebook or Jupyter Lab installed in your Python environment, please follow these instructions. Now open a terminal on your computer and start a notebook/lab session. A file-viewer will open in an internet browser; pick a directory where you are okay with saving some PyTorch model weights. Create a notebook called Building-Workflows.ipynb. You can then follow along with this tutorial by copying, pasting, and running the code blocks below in the cells of your notebook.

Note: you may also need to install the ipywidgets package in your Python environment to configure the notebook to display ipywidgets:

$ pip install ipywidgets

Imports#

[1]:

from pathlib import Path
from typing import Optional, Tuple, Union

import matplotlib.pyplot as plt
import torch as tr

[2]:

# Hydra and hydra-zen
from hydra.core.config_store import ConfigStore
from hydra_zen import MISSING, builds, instantiate, load_from_yaml, make_config

# Lightning
from pytorch_lightning import LightningModule, Trainer

# sklearn and torch
from sklearn.model_selection import StratifiedKFold
from torch import Tensor, nn
from torch.optim import Optimizer
from torch.utils.data import DataLoader, Subset
from torchmetrics import Accuracy
from torchvision import transforms
from torchvision.datasets import MNIST

# rAI-toolbox
from rai_toolbox._typing import Partial
from rai_toolbox.mushin import load_from_checkpoint
from rai_toolbox.mushin.lightning import MetricsCallback
from rai_toolbox.mushin.workflows import (
    MultiRunMetricsWorkflow,
    RobustnessCurve,
    multirun,
)

from rai_toolbox.optim import L2ProjectedOptim, LinfProjectedOptim
from rai_toolbox.perturbations import gradient_ascent

Experiment Functions and Classes#

Here we define two Neural Network models, a fully linear neural network and a convolutional neural network.

[3]:

class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Flatten(1),
            nn.Linear(28 * 28, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 10),
        )

    def forward(self, x):
        return self.model(x)

class ConvModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Conv2d(1, 32, 5, padding="same"),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(3),
            nn.Conv2d(32, 32, 3, padding="same"),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(3),
            nn.Conv2d(32, 32, 3, padding="same"),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 10, 3),
            nn.Flatten(1),
        )

    def forward(self, x):
        return self.model(x)

Next lets define a function that takes the MNIST dataset and splits the data into training and validation sets using SciKit-Learn’s StratifiedKFold. This allows us to split the dataset into “folds” and select the fold for each experiment.

[4]:

def split_dataset(
    dataset: MNIST, n_splits: int, fold: int, random_state: int = 49
) -> Tuple[Subset, Subset]:
    """Provide training and validation splits using `sklearn.model_selection.StratifiedKfold`"""

    kfold = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=random_state)
    train_indices, val_indices = list(
        kfold.split(range(len(dataset)), dataset.targets)
    )[fold]
    return Subset(dataset, train_indices), Subset(dataset, val_indices)

Now define the LightningModule for training and testing. This describes how we:

Load our data
Process a batch of data with our model (both with and without adversarial perturbations)
Update our model’s parameters during training

Note that we specifically design this lightning module to log the following metrics:

Loss and accuracy for cross-validation training
Adversarial loss, adversarial accuracy, and clean accuracy for robustness analysis

These metrics will be saved during each of our runs, and we will load and aggregate these metrics to analyze our results.

[5]:

class StandardModule(LightningModule):
    def __init__(
        self,
        *,
        model: nn.Module,
        dataset: MNIST,
        optimizer: Optional[Partial[Optimizer]] = None,
        perturber=None,
        fold: int = 0,
        n_splits: int = 5,
        batch_size: int = 100,
        num_workers: int = 4,
    ) -> None:
        super().__init__()
        self.dataset = dataset
        self.optimizer = optimizer
        self.criterion = nn.CrossEntropyLoss()
        self.model = model
        self.perturber = perturber
        self.n_splits = n_splits
        self.fold = fold
        self.batch_size = batch_size
        self.num_workers = num_workers

        # Metrics
        self.acc_metric = Accuracy(task="multiclass", num_classes=10)
        if self.perturber:
            self.clean_acc_metric = Accuracy(task="multiclass", num_classes=10)

    def forward(self, data: Tensor) -> Tensor:
        return self.model(data)

    def train_dataloader(self) -> DataLoader:
        train_dataset, _ = split_dataset(self.dataset, self.n_splits, self.fold)
        return DataLoader(
            train_dataset,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            shuffle=True,
        )

    def val_dataloader(self) -> DataLoader:
        _, val_dataset = split_dataset(self.dataset, self.n_splits, self.fold)
        return DataLoader(
            val_dataset, batch_size=self.batch_size, num_workers=self.num_workers
        )

    def test_dataloader(self) -> DataLoader:
        return DataLoader(
            self.dataset, batch_size=self.batch_size, num_workers=self.num_workers
        )

    def configure_optimizers(self) -> Optional[Optimizer]:
        if self.optimizer:
            return self.optimizer(self.model.parameters())
        return None

    def _step(self, batch, stage: str) -> Tensor:
        data_orig, target = batch

        if self.perturber:
            with tr.no_grad():
                output = self.model(data_orig)
                loss = self.criterion(output, target)
                acc = self.clean_acc_metric(output, target)
                self.log(f"{stage}_clean_accuracy", acc)

            inference_tensors = tr.is_inference_mode_enabled()
            with tr.inference_mode(mode=False), tr.enable_grad():
                if inference_tensors:
                    # we need to clone in order to support grad mode
                    data_orig = data_orig.clone()
                    target = target.clone()

                data, adv_loss = self.perturber(
                    model=self.model, data=data_orig, target=target
                )
            self.log(f"{stage}_adversarial_loss", adv_loss.mean().item())

        else:
            data = data_orig

        output = self.model(data)
        loss = self.criterion(output, target)
        acc = self.acc_metric(output, target)
        self.log(f"{stage}_loss", loss)
        self.log(f"{stage}_accuracy", acc)
        return loss

    def training_step(self, batch, batch_idx) -> Tensor:
        return self._step(batch, "train")

    def validation_step(self, batch, batch_idx) -> Tensor:
        return self._step(batch, "val")

    def test_step(self, batch, batch_idx) -> Tensor:
        return self._step(batch, "test")

Configuring our experiments with hydra-zen#

Now we use hydra-zen to create “configs” for all of the components of our experiments. Each config describes an interface and/or object in our experiment that we want to be able to modify from run to run. They will also serve to make our work self-documenting and reproducible.

[6]:

Augmentations = builds(
    transforms.Compose,
    [builds(transforms.RandomCrop, size=28, padding=4), builds(transforms.ToTensor)],
)
TrainDataset = builds(
    MNIST, root="${data_dir}", train=True, transform=Augmentations, download=True
)
TestDataset = builds(
    MNIST,
    root="${data_dir}",
    train=False,
    transform=builds(transforms.ToTensor),
    download=True,
)
ConvModelCfg = builds(ConvModel)
LinearModelCfg = builds(LinearModel)
Optim = builds(tr.optim.SGD, lr=0.1, zen_partial=True)


L2PGD = builds(L2ProjectedOptim, zen_partial=True)
LinfPGD = builds(LinfProjectedOptim, zen_partial=True)


def lr_for_pgd(epsilon, num_steps):
    return 2.5 * epsilon / num_steps


Perturber = builds(
    gradient_ascent,
    optimizer="${optimizer}",
    epsilon="${epsilon}",
    steps="${steps}",
    lr=builds(lr_for_pgd, "${epsilon}", "${steps}"),
    zen_partial=True,
    populate_full_signature=True,
)

PLModule = builds(
    StandardModule,
    model="${model}",
    fold="${fold}",
    n_splits="${n_splits}",
    dataset=TrainDataset,
    optimizer=Optim,
    perturber="${perturber}",
    populate_full_signature=True,
)


EvalPLModule = builds(
    StandardModule,
    model="${model}",
    dataset=TestDataset,
    perturber="${perturber}",
    populate_full_signature=True,
)

We configure our trainer to use MetricsCallback, which will instruct PyTorch Lightning to automatically save our logged metrics as a dictionary in a file named “fit_metrics.pt” and “test_metrics.pt” for training and evaluation, respectively.

[7]:

TrainerCfg = builds(
    Trainer,
    max_epochs=10,
    accelerator="auto",
    devices=1,
    enable_progress_bar=False,
    enable_model_summary=False,
    callbacks=[builds(MetricsCallback)],
    populate_full_signature=True,
)

Now we use Hydra’s ConfigStore API to create named configuration groups that can be specified/swapped when we run our workflow. Let’s make it easy to swap both models and optimizers by-name.

[8]:

cs = ConfigStore.instance()
cs.store(name="cnn", group="model", node=ConvModelCfg)
cs.store(name="linear", group="model", node=LinearModelCfg)
cs.store(name="l2pgd", group="optimizer", node=L2PGD)
cs.store(name="linfpgd", group="optimizer", node=LinfPGD)

Building Workflows for Cross Validation Training and Adversarial Robustness Analysis#

Getting Started#

Installing `rai_toolbox`#

Opening a Jupyter notebook#

Imports#

Experiment Functions and Classes#

Configuring our experiments with hydra-zen#

Cross Validation Workflow#

Robustness Curve#

Building Workflows for Cross Validation Training and Adversarial Robustness Analysis#

Getting Started#

Installing rai_toolbox#

Opening a Jupyter notebook#

Imports#

Experiment Functions and Classes#

Configuring our experiments with hydra-zen#

Cross Validation Workflow#

Robustness Curve#

Installing `rai_toolbox`#