How to measure dataset sufficiency for image classification¶

import os
import random
from collections.abc import Sequence
from typing import cast

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchmetrics
from maite_datasets.image_classification import MNIST
from tabulate import tabulate
from torch.utils.data import DataLoader, Dataset, Subset

from dataeval.data import Select
from dataeval.data.selections import Limit
from dataeval.workflows import Sufficiency

np.random.seed(0)
np.set_printoptions(formatter={"float": lambda x: f"{x:0.4f}"})
torch.manual_seed(0)
torch.set_float32_matmul_precision("high")
device = "cuda" if torch.cuda.is_available() else "cpu"
torch._dynamo.config.suppress_errors = True

random.seed(0)
torch.use_deterministic_algorithms(True)
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

# Configure the dataset transforms

transforms = [
    lambda x: x / 255.0,  # scale to [0, 1]
    lambda x: x.astype(np.float32),  # convert to float32
]

# Download the mnist dataset and apply the transforms and subset the data

train_ds = Select(MNIST(root="./data", image_set="train", transforms=transforms,download=True),selections=[Limit(2500)])  # fmt: skip # noqa: E501
test_ds = Select(MNIST(root="./data", image_set="test", transforms=transforms, download=True), selections=[Limit(500)])

# Define our network architecture
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(6400, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = torch.flatten(x, 1)  # flatten all dimensions except batch
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


# Compile the model
model = torch.compile(Net().to(device))

# Type cast the model back to Net as torch.compile returns a Unknown
# Nothing internally changes from the cast; we are simply signaling the type
model = cast(Net, model)

def custom_train(model: nn.Module, dataset: Dataset, indices: Sequence[int]):
    # Defined only for this testing scenario
    criterion = torch.nn.CrossEntropyLoss().to(device)
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
    epochs = 10

    # Define the dataloader for training
    dataloader = DataLoader(Subset(dataset, indices), batch_size=8)

    for epoch in range(epochs):
        for batch in dataloader:
            # Load data/images to device
            X = torch.Tensor(batch[0]).to(device)
            # Load one-hot encoded targets/labels to device
            y = torch.argmax(torch.asarray(batch[1], dtype=torch.int).to(device), dim=1)
            # Zero out gradients
            optimizer.zero_grad()
            # Forward propagation
            outputs = model(X)
            # Compute loss
            loss = criterion(outputs, y)
            # Back prop
            loss.backward()
            # Update weights/parameters
            optimizer.step()

def custom_eval(model: nn.Module, dataset: Dataset) -> dict[str, float]:
    # Metrics of interest
    metrics = {
        "Accuracy": torchmetrics.Accuracy(task="multiclass", num_classes=10).to(device),
        "AUROC": torchmetrics.AUROC(task="multiclass", num_classes=10).to(device),
        "TPR at 0.5 Fixed FPR": torchmetrics.ROC(task="multiclass", average="macro", num_classes=10).to(device),
    }
    result = {}
    # Set model layers into evaluation mode
    model.eval()
    dataloader = DataLoader(dataset, batch_size=8)
    # Tell PyTorch to not track gradients, greatly speeds up processing
    with torch.no_grad():
        for batch in dataloader:
            # Load data/images to device
            X = torch.Tensor(batch[0]).to(device)
            # Load one-hot encoded targets/labels to device
            y = torch.argmax(torch.asarray(batch[1], dtype=torch.int).to(device), dim=1)
            preds = model(X)
            for metric in metrics.values():
                metric.update(preds, y)
        # Compute ROC curve
        false_positive_rate, true_positive_rate, _ = metrics["TPR at 0.5 Fixed FPR"].compute()
        # determine interval to examine
        desired_rate = 0.5
        closest_desired_index = torch.argmin(torch.abs(false_positive_rate - desired_rate)).item()
        # return corresponding tpr value
        result["TPR at 0.5 Fixed FPR"] = true_positive_rate[closest_desired_index].cpu()
        result["Accuracy"] = metrics["Accuracy"].compute().cpu()
        result["AUROC"] = metrics["AUROC"].compute().cpu()
    return result

# Instantiate sufficiency metric
suff = Sufficiency(
    model=model,  # type: ignore
    train_ds=train_ds,  # type: ignore
    test_ds=test_ds,  # type: ignore
    train_fn=custom_train,
    eval_fn=custom_eval,
    runs=5,
    substeps=10,
)

# Train & test model
output = suff.evaluate()

# Print out sufficiency output in a table format
formatted = {"Steps": output.steps, **output.averaged_measures}
print(tabulate(formatted, headers=list(formatted), tablefmt="pretty"))

# Print out projected output values
projection = output.project([1000, 2500, 5000])
projected = {"Steps": projection.steps, **projection.averaged_measures}
print(tabulate(projected, list(projected), tablefmt="pretty"))

# Plot the output using the convenience function
_ = output.plot()

# Initialize the array of desired thresholds to apply to all metrics
desired_values = np.array([0.90, 0.93, 0.99])
metrics = ["Accuracy", "AUROC", "TPR at 0.5 Fixed FPR"]
evaluated_metrics = {}

for metric in metrics:
    evaluated_metrics[metric] = desired_values
# Evaluate the learning curve to infer the needed amount of training data
samples_needed = output.inv_project(evaluated_metrics)

# Print the amount of needed data needed to achieve the thresholds
for metric, samples in samples_needed.items():
    print(f"{metric}")
    for index, sample_size in enumerate(samples):
        print(
            f"To achieve {int(evaluated_metrics[metric][index] * 100)}% {metric},"
            f" {int(sample_size)} samples are needed."
        )
    print()