API Reference: Models

Base model

auditml.models.base.BaseModel

Bases: ABC, Module

Base class that all AuditML target/shadow models must extend.

Subclasses must implement forward and get_features.

Source code in src/auditml/models/base.py

class BaseModel(ABC, nn.Module):
    """Base class that all AuditML target/shadow models must extend.

    Subclasses must implement ``forward`` and ``get_features``.
    """

    @abstractmethod
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Full forward pass returning logits of shape ``(batch, num_classes)``."""
        ...

    @abstractmethod
    def get_features(self, x: torch.Tensor) -> torch.Tensor:
        """Return the penultimate-layer feature vector.

        Used by membership inference attacks to analyse the model's internal
        representations. Shape: ``(batch, feature_dim)``.
        """
        ...

forward(x: torch.Tensor) -> torch.Tensor abstractmethod

Full forward pass returning logits of shape (batch, num_classes).

Source code in src/auditml/models/base.py

@abstractmethod
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Full forward pass returning logits of shape ``(batch, num_classes)``."""
    ...

get_features(x: torch.Tensor) -> torch.Tensor abstractmethod

Return the penultimate-layer feature vector.

Used by membership inference attacks to analyse the model's internal representations. Shape: (batch, feature_dim).

Source code in src/auditml/models/base.py

@abstractmethod
def get_features(self, x: torch.Tensor) -> torch.Tensor:
    """Return the penultimate-layer feature vector.

    Used by membership inference attacks to analyse the model's internal
    representations. Shape: ``(batch, feature_dim)``.
    """
    ...

CNN architectures

auditml.models.cnn

Simple CNN architecture for MNIST / CIFAR-10 / CIFAR-100.

The architecture is intentionally compact — two conv blocks followed by two fully-connected layers. This mirrors common MIA literature baselines and trains quickly, which matters when training multiple shadow models.

SimpleCNN

Bases: BaseModel

Configurable 2-block CNN.

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (1 for MNIST, 3 for CIFAR).	1
num_classes	int	Number of output classes.	10
input_size	int	Spatial dimension of the input (28 for MNIST, 32 for CIFAR).	28
feature_dim	int	Width of the penultimate fully-connected layer.	128

Source code in src/auditml/models/cnn.py

class SimpleCNN(BaseModel):
    """Configurable 2-block CNN.

    Parameters
    ----------
    input_channels:
        Number of input channels (1 for MNIST, 3 for CIFAR).
    num_classes:
        Number of output classes.
    input_size:
        Spatial dimension of the input (28 for MNIST, 32 for CIFAR).
    feature_dim:
        Width of the penultimate fully-connected layer.
    """

    def __init__(
        self,
        input_channels: int = 1,
        num_classes: int = 10,
        input_size: int = 28,
        feature_dim: int = 128,
    ) -> None:
        super().__init__()

        self.conv1 = nn.Conv2d(input_channels, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.dropout_conv = nn.Dropout2d(0.25)
        self.dropout_fc = nn.Dropout(0.5)

        # After two pools: spatial size = input_size // 4
        flat_size = 64 * (input_size // 4) ** 2

        self.fc1 = nn.Linear(flat_size, feature_dim)
        self.fc2 = nn.Linear(feature_dim, num_classes)

        self._feature_dim = feature_dim

    # ----- forward / features --------------------------------------------

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self._backbone(x)
        x = self.fc2(x)
        return x

    def get_features(self, x: torch.Tensor) -> torch.Tensor:
        return self._backbone(x)

    # ----- internals -----------------------------------------------------

    def _backbone(self, x: torch.Tensor) -> torch.Tensor:
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.dropout_conv(x)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = self.dropout_fc(x)
        return x

ResNet

auditml.models.resnet

Small ResNet variant for CIFAR datasets.

A compact ResNet-18-style architecture adapted for 32 × 32 inputs (no aggressive stem downsampling). Achieves stronger accuracy than SimpleCNN on CIFAR-10/100, useful for testing whether a more capable model leaks more privacy.

SmallResNet

Bases: BaseModel

ResNet-18-style network for 32×32 inputs.

Uses GroupNorm instead of BatchNorm so the model is compatible with Opacus (differential-privacy training) out of the box.

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (1 for MNIST, 3 for CIFAR).	3
num_classes	int	Number of output classes.	10
feature_dim	int	Width of the penultimate layer (default 512).	512

Source code in src/auditml/models/resnet.py

class SmallResNet(BaseModel):
    """ResNet-18-style network for 32×32 inputs.

    Uses GroupNorm instead of BatchNorm so the model is compatible with
    Opacus (differential-privacy training) out of the box.

    Parameters
    ----------
    input_channels:
        Number of input channels (1 for MNIST, 3 for CIFAR).
    num_classes:
        Number of output classes.
    feature_dim:
        Width of the penultimate layer (default 512).
    """

    def __init__(
        self,
        input_channels: int = 3,
        num_classes: int = 10,
        feature_dim: int = 512,
    ) -> None:
        super().__init__()
        self.in_planes = 64

        self.conv1 = nn.Conv2d(
            input_channels, 64, kernel_size=3, stride=1, padding=1, bias=False,
        )
        self.bn1 = nn.GroupNorm(32, 64)

        self.layer1 = self._make_layer(64, 2, stride=1)
        self.layer2 = self._make_layer(128, 2, stride=2)
        self.layer3 = self._make_layer(256, 2, stride=2)
        self.layer4 = self._make_layer(feature_dim, 2, stride=2)

        self.fc = nn.Linear(feature_dim, num_classes)
        self._feature_dim = feature_dim

    # ----- forward / features --------------------------------------------

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        feat = self._backbone(x)
        return self.fc(feat)

    def get_features(self, x: torch.Tensor) -> torch.Tensor:
        return self._backbone(x)

    # ----- internals -----------------------------------------------------

    def _make_layer(self, planes: int, num_blocks: int, stride: int) -> nn.Sequential:
        strides = [stride] + [1] * (num_blocks - 1)
        layers: list[nn.Module] = []
        for s in strides:
            layers.append(_BasicBlock(self.in_planes, planes, s))
            self.in_planes = planes
        return nn.Sequential(*layers)

    def _backbone(self, x: torch.Tensor) -> torch.Tensor:
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = F.adaptive_avg_pool2d(out, 1)
        return torch.flatten(out, 1)