Certified Training with CTRAIN¶

In this example, we train the CNN7 architecture proposed by Shi et al. on MNIST using IBP with Shi-style improvements. We want the network to be certifiably robust against ℓ∞ perturbations with radius ε = 0.1. The example uses the real MNIST dataset and a 70-epoch training run; a CUDA GPU is recommended.

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 and define the model. The loader downloads missing data to ./data and supplies the normalisation metadata used during training. The radius remains specified in raw pixel units.

In [ ]:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
input_shape = (1, 28, 28)
train_loader, test_loader = load_mnist(batch_size=128, val_split=False, data_root="./data")
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, test_loader = load_mnist(batch_size=128, val_split=False, data_root="./data") model = CNN7_Shi(in_shape=input_shape, n_classes=10)

To train the network certifiably, we wrap it with ShiIBPModelWrapper. We configure one warm-up epoch on clean loss and 40 ramp-up epochs during which the perturbation radius increases to its target. The learning rate decreases at epochs 50 and 60.

To use another training method, import its wrapper and supply the corresponding method parameters.

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, )

We initiate training by calling train_model on the wrapped model.

In [ ]:

wrapped_model.train_model(train_loader)

wrapped_model.train_model(train_loader)

Finally, we save the resulting model weights. The model-evaluation notebook loads this file using the same architecture and wrapper.

In [ ]:

torch.save(wrapped_model.state_dict(), "mnist_ibp.pt")

torch.save(wrapped_model.state_dict(), "mnist_ibp.pt")