aiml.surrogate_model.models — Adversarial Insight ML documentation

# Source code for aiml.surrogate_model.models

    """
    models.py

    This module contains utility functions and PyTorch Lightning modules for
    working with the CIFAR-10 dataset. The VGG16 BN model is used as a
    substitute for the black box model. This functions and classes in this
    file are used in the "create_surrogate_model.py" file.
    """

    import pytorch_lightning as pl
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    import torch.optim as optim
    import torchvision as tv
    from torchmetrics import Accuracy

    def create_substitute_model(num_classes, num_channels):
        """
        Create a substitute model based on the input model.

        Parameters:
            num_classes (int): The number of output classes for the model.
            num_channels (int): The number of input channels.

        Returns:
            nn.Module: The created substitute model.
        """
        model = tv.models.vgg16(weights=None, num_classes=num_classes)
        model.features[0] = nn.Conv2d(
            num_channels, 64, kernel_size=3, stride=1, padding=1, bias=False)
        model.features[4] = nn.Identity()
        return model

    class Surrogate(pl.LightningModule):
        """A PyTorch Lightning module representing a surrogate model.

        This surrogate model is designed to mimic the behavior of an oracle model.

        Attributes:
            oracle (nn.Module): The oracle model for reference.
            substitute (nn.Module): The surrogate model to be trained.
            loss_fn (Callable): The loss function for surrogate training.
            accuracy (Accuracy): A metric for computing accuracy during training/validation.
        """

        def __init__(self, lr, num_training_batches, oracle, substitute, loss_fn, num_classes, softmax=True):
            super().__init__()
            self.save_hyperparameters(ignore=["oracle", "substitute", "loss_fn"])

            self.oracle = oracle
            for param in self.oracle.parameters():
                param.requires_grad = False
            self.substitute = substitute
            self.loss_fn = loss_fn
            self.accuracy = Accuracy(task="multiclass", num_classes=num_classes)

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

        def training_step(self, batch, batch_idx):
            x, _ = batch
            out = self(x)
            if self.hparams.softmax:
                out = F.log_softmax(out, 1)
            self.oracle.eval()
            with torch.no_grad():
                out_oracle = self.oracle(x)
            loss = self.loss_fn(out, out_oracle)
            self.log(f"train_loss", loss)
            return loss

        def validation_step(self, batch, batch_idx):
            x, y = batch
            out = self(x)
            if self.hparams.softmax:
                out = F.log_softmax(out, 1)
            out_oracle = self.oracle(x)
            loss = self.loss_fn(out, out_oracle)
            pred_oracle = out_oracle.argmax(1)
            acc = self.accuracy(out, y)
            match_oracle = self.accuracy(out, pred_oracle)
            self.log(f"val_loss", loss)
            self.log(f"val_acc", acc)
            self.log(f"val_match", match_oracle)

        def predict_step(
            self, batch, batch_idx, dataloader_idx=0
        ):
            x, _ = batch
            out = self(x)
            return out

        def configure_optimizers(self):
            optimizer = torch.optim.SGD(
                self.substitute.parameters(),
                lr=self.hparams.lr,
                momentum=0.9,
                weight_decay=5e-4,
            )
            scheduler_dict = {
                "scheduler": optim.lr_scheduler.OneCycleLR(
                    optimizer,
                    epochs=self.trainer.max_epochs,
                    steps_per_epoch=self.hparams.num_training_batches,
                    max_lr=0.1,
                ),
                "interval": "step",
            }
            return {"optimizer": optimizer, "lr_scheduler": scheduler_dict}

    class LogSoftmaxModule(pl.LightningModule):
        """A PyTorch Lightning module that wraps a model and applies LogSoftmax to its output.

        This module is designed to enhance the functionality of an existing neural network model
        by applying LogSoftmax to its output. It can be used for various machine learning tasks
        such as classification.

        Attributes:
            model (nn.Module): The underlying model to wrap with LogSoftmax.
        """

        def __init__(self, model):
            super().__init__()
            self.model = model

        def forward(self, x):
            x = self.model(x)
            x = F.log_softmax(x, dim=1)
            return x

        def predict_step(
            self, batch, batch_idx, dataloader_idx=0
        ):
            x, _ = batch
            out = self(x)
            return out