Hyperparameter Optimisation with CTRAIN¶

In this example, we use CTRAIN to optimise natural and certified validation accuracy on MNIST. The search returns a Pareto front of configurations representing different trade-offs between the two objectives.

First, we import PyTorch and the required functions from CTRAIN.

In [ ]:

import torch

from CTRAIN.data_loaders import load_mnist
from CTRAIN.model_definitions import CNN7_Shi
from CTRAIN.model_wrappers import ShiIBPModelWrapper

import torch from CTRAIN.data_loaders import load_mnist from CTRAIN.model_definitions import CNN7_Shi from CTRAIN.model_wrappers import ShiIBPModelWrapper

Now, we load MNIST with a validation holdout and construct the model. We reserve the test set for assessing the selected model after the search.

In [ ]:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
input_shape = (1, 28, 28)
train_loader, val_loader, test_loader = load_mnist(batch_size=128, data_root="./data", seed=0)
model = CNN7_Shi(in_shape=input_shape, n_classes=10)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") input_shape = (1, 28, 28) train_loader, val_loader, test_loader = load_mnist(batch_size=128, data_root="./data", seed=0) model = CNN7_Shi(in_shape=input_shape, n_classes=10)

We wrap the network with ShiIBPModelWrapper. Each trial uses the same architecture, evaluation radius and 70-epoch budget. CTRAIN chooses the remaining training parameters from the method-specific configuration space.

In [ ]:

wrapped_model = ShiIBPModelWrapper(
    model=model,
    input_shape=input_shape,
    eps=0.1,
    num_epochs=70,
    warm_up_epochs=1,
    ramp_up_epochs=40,
    lr_decay_kwargs={"milestones": [50, 60], "gamma": 0.2},
    device=device,
)

wrapped_model = ShiIBPModelWrapper( model=model, input_shape=input_shape, eps=0.1, num_epochs=70, warm_up_epochs=1, ramp_up_epochs=40, lr_decay_kwargs={"milestones": [50, 60], "gamma": 0.2}, device=device, )

Now, we call hpo. Here, we use three trials and Optuna's NSGA-II sampler to demonstrate the interface. A larger search can use the default BoTorch sampler and a larger trial or time budget. Full MNIST trials still require substantial training time.

In [ ]:

pareto_front = wrapped_model.hpo(
    train_loader=train_loader,
    val_loader=val_loader,
    budget_trials=3,
    eval_samples=1000,
    sampler="nsgaii",
    seed=0,
    output_dir="./hpo/mnist_ibp",
)

for trial in pareto_front:
    print(trial["metrics"])

pareto_front = wrapped_model.hpo( train_loader=train_loader, val_loader=val_loader, budget_trials=3, eval_samples=1000, sampler="nsgaii", seed=0, output_dir="./hpo/mnist_ibp", ) for trial in pareto_front: print(trial["metrics"])

Finally, we select the model with the highest certified validation accuracy, load its weights and evaluate it on the held-out test set. Other applications may prefer a different point on the front, for example one satisfying a natural-accuracy requirement.

In [ ]:

selected_trial = max(pareto_front, key=lambda trial: trial["metrics"]["cert_acc"])
wrapped_model.load_state_dict(
    torch.load(selected_trial["checkpoint_path"], map_location=device, weights_only=True),
)
wrapped_model.eval()
std_acc, cert_acc, adv_acc = wrapped_model.evaluate(test_loader, test_samples=1000)
print(f"Natural: {std_acc:.2%}, certified: {cert_acc:.2%}, adversarial: {adv_acc:.2%}")

selected_trial = max(pareto_front, key=lambda trial: trial["metrics"]["cert_acc"]) wrapped_model.load_state_dict( torch.load(selected_trial["checkpoint_path"], map_location=device, weights_only=True), ) wrapped_model.eval() std_acc, cert_acc, adv_acc = wrapped_model.evaluate(test_loader, test_samples=1000) print(f"Natural: {std_acc:.2%}, certified: {cert_acc:.2%}, adversarial: {adv_acc:.2%}")

The study and trial weights are saved in the output directory. Reusing the same compatible study continues towards the total budget_trials; it does not request that many additional trials. The HPO guide explains constraints, study continuation and scalar objectives.