Skip to content

SABR

Model-wrapper interface for SABR training and robustness evaluation.

SABRModelWrapper

Bases: CTRAINWrapper

Wrap a PyTorch classifier for SABR certified training.

This wrapper configures SABR training, using IBP on adversarially selected input subregions as proposed by Müller et al. (2023). Training is invoked through train_model; evaluation and hyperparameter optimisation use the common CTRAINWrapper interface. The training loop, loss computation and bound propagation are implemented in the corresponding library functions.

See https://arxiv.org/abs/2210.04871.

Source code in CTRAIN/model_wrappers/certified/sabr_model_wrapper.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
class SABRModelWrapper(CTRAINWrapper):
    """Wrap a PyTorch classifier for SABR certified training.

    This wrapper configures SABR training, using IBP on adversarially selected input subregions as proposed by Müller et al. (2023).
    Training is invoked through train_model; evaluation and hyperparameter
    optimisation use the common CTRAINWrapper interface. The training loop,
    loss computation and bound propagation are implemented in the corresponding
    library functions.

    See https://arxiv.org/abs/2210.04871.
    """

    def __init__(self, model, input_shape, eps, num_epochs, train_eps_factor=1, optimizer_func=torch.optim.Adam, lr=0.0005, warm_up_epochs=1, ramp_up_epochs=70,
                 lr_scheduler_func=torch.optim.lr_scheduler.MultiStepLR, lr_decay_kwargs=dict(milestones=(80, 90), gamma=0.2), gradient_clip=10, l1_reg_weight=0.000001,
                 shi_reg_weight=.5, shi_reg_decay=True, sabr_subselection_ratio=.2, pgd_steps=8,
                 pgd_alpha=0.5, pgd_restarts=1, pgd_early_stopping=False, pgd_alpha_decay_factor=.1,
                 pgd_decay_milestones=(4,7), pgd_eps_factor=1, checkpoint_save_path=None, checkpoint_save_interval=10,
                 bound_opts=dict(conv_mode='patches', relu='adaptive'), device=torch.device('cuda'), population_bn=False, relu_upper_retention=1.0):
        """Configure the model, training schedule and SABR hyperparameters.

        Args:
            model (torch.nn.Module): The model to be trained.
            input_shape (tuple): Shape of the input data.
            eps (float): Epsilon value describing the perturbation the network should be
                certifiably robust against.
            num_epochs (int): Number of epochs for training.
            train_eps_factor (float): Factor for training epsilon.
            optimizer_func (torch.optim.Optimizer): Optimizer function.
            lr (float): Learning rate.
            warm_up_epochs (int): Number of warm-up epochs, i.e. epochs where the model is
                trained on clean loss.
            ramp_up_epochs (int): Number of ramp-up epochs, i.e. epochs where the epsilon is
                gradually increased to the target train epsilon.
            lr_scheduler_func (callable): Constructor for the learning-rate scheduler.
            lr_decay_kwargs (dict): Scheduler arguments, such as milestones and gamma for MultiStepLR.
            gradient_clip (float): Gradient clipping value.
            l1_reg_weight (float): L1 regularisation weight.
            shi_reg_weight (float): SHI regularisation weight.
            shi_reg_decay (bool): Whether to decay SHI regularisation during the ramp up phase.
            sabr_subselection_ratio (float): Ratio of subselection for SABR (lambda in original
                publication).
            pgd_steps (int): Number of PGD steps for adversarial loss computation.
            pgd_alpha (float): PGD step size for adversarial loss calculation.
            pgd_restarts (int): Number of PGD restarts for adversarial loss calculation.
            pgd_early_stopping (bool): Whether to use early stopping in PGD during adversarial
                loss calculation.
            pgd_alpha_decay_factor (float): PGD alpha decay factor.
            pgd_decay_milestones (tuple): Milestones for PGD alpha decay.
            pgd_eps_factor (float): Factor for epsilon used during PGD attack.
            checkpoint_save_path (str): Path to save checkpoints.
            checkpoint_save_interval (int): Interval for saving checkpoints.
            bound_opts (dict): Options for bounding according to the auto_LiRPA documentation.
            device (torch.device): Device to run the training on.
            population_bn (bool): Recalibrate BatchNorm on training data after every epoch.
            relu_upper_retention (float): Fraction of unstable ReLU upper bounds retained during
                training; default 1.
        """
        validate_hybrid_options(relu_upper_retention, pgd_eps_factor)
        super().__init__(model, eps, input_shape, train_eps_factor, lr, optimizer_func, lr_scheduler_func, lr_decay_kwargs, bound_opts, device, checkpoint_save_path=checkpoint_save_path, checkpoint_save_interval=checkpoint_save_interval, population_bn=population_bn)
        self.relu_upper_retention = relu_upper_retention
        self.cert_train_method = 'sabr'
        self.num_epochs = num_epochs
        self.lr = lr
        self.warm_up_epochs = warm_up_epochs
        self.ramp_up_epochs = ramp_up_epochs
        self.gradient_clip = gradient_clip
        self.l1_reg_weight = l1_reg_weight
        self.shi_reg_weight = shi_reg_weight
        self.shi_reg_decay = shi_reg_decay
        self.optimizer_func = optimizer_func
        self.sabr_subselection_ratio = sabr_subselection_ratio
        self.pgd_steps = pgd_steps
        self.pgd_alpha = pgd_alpha
        self.pgd_restarts = pgd_restarts
        self.pgd_early_stopping = pgd_early_stopping
        self.pgd_alpha_decay_factor = pgd_alpha_decay_factor
        self.pgd_decay_milestones = pgd_decay_milestones
        self.pgd_eps_factor = pgd_eps_factor

    def train_model(self, train_loader, val_loader=None, start_epoch=0, end_epoch=None):
        """Train the wrapped model using the configured method and schedule.

        The wrapper forwards its configuration to the corresponding library trainer.
        start_epoch and end_epoch select a segment of the full schedule, retaining
        its original warm-up, ramp-up and learning-rate milestones.

        Args:
            train_loader (torch.utils.data.DataLoader): DataLoader for training data.
            val_loader (torch.utils.data.DataLoader, optional): DataLoader for validation data.
            start_epoch (int, optional): Epoch to start training from. Initialises learning rate
                and epsilon schedulers accordingly.
            end_epoch (int, optional): Epoch to prematurely end training at.

        Returns:
            (auto_LiRPA.BoundedModule): Trained model.
        """
        eps_std = self.train_eps / train_loader.std if train_loader.normalised else torch.tensor(self.train_eps)
        eps_std = torch.reshape(eps_std, (*eps_std.shape, 1, 1))
        trained_model = sabr_train_model(
            original_model=self.original_model,
            hardened_model=self.bounded_model,
            train_loader=train_loader,
            val_loader=val_loader,
            start_epoch=start_epoch,
            end_epoch=end_epoch,
            num_epochs=self.num_epochs,
            eps=self.train_eps,
            eps_std=eps_std,
            eps_schedule=(self.warm_up_epochs, self.ramp_up_epochs),
            eps_scheduler_args={},
            optimizer=self.optimizer,
            lr_scheduler=self.lr_scheduler,
            n_classes=self.n_classes,
            gradient_clip=self.gradient_clip,
            l1_regularisation_weight=self.l1_reg_weight,
            shi_regularisation_weight=self.shi_reg_weight,
            shi_reg_decay=self.shi_reg_decay,
            eps_schedule_unit='epoch',
            lr_decay_schedule_unit='epoch',
            subselection_ratio=self.sabr_subselection_ratio,
            pgd_steps=self.pgd_steps,
            pgd_step_size=self.pgd_alpha,
            pgd_restarts=self.pgd_restarts,
            pgd_early_stopping=self.pgd_early_stopping,
            pgd_decay_checkpoints=self.pgd_decay_milestones,
            pgd_decay_factor=self.pgd_alpha_decay_factor,
            pgd_eps_factor=self.pgd_eps_factor,
            results_path=self.checkpoint_path,
            checkpoint_save_interval=self.checkpoint_save_interval,
            population_bn=self.population_bn,
            relu_upper_retention=self.relu_upper_retention,
            device=self.device, selection_state=self.selection_state, history=self.history, validation_eps=self.eps)

        self.epoch = self.num_epochs if end_epoch is None else end_epoch
        return trained_model

    def _hpo_runner(self, config, seed, epochs, train_loader, val_loader, output_dir, cert_eval_samples=1000, nat_loss_weight=1., adv_loss_weight=1., cert_loss_weight=1., complete_verify=False):
        """Train and evaluate one hyperparameter configuration.

        Construct a method-specific wrapper, train it and save its inference weights.
        Return the negative weighted accuracy and the natural, certified and
        adversarial validation accuracies for the optimisation procedure.

        Args:
            config (dict): Configuration of hyperparameters.
            seed (int): Seed used.
            epochs (int): Number of epochs for training.
            train_loader (torch.utils.data.DataLoader): DataLoader for training data.
            val_loader (torch.utils.data.DataLoader): DataLoader for validation data.
            output_dir (str): Directory to save output.
            cert_eval_samples (int, optional): Number of samples for certification evaluation.
            nat_loss_weight (float, optional): Weight for the natural accuracy in the loss
                function.
            adv_loss_weight (float, optional): Weight for the adversarial accuracy in the loss
                function.
            cert_loss_weight (float, optional): Weight for the certified accuracy in the loss
                function.
            complete_verify (bool): Whether to evaluate the trial through complete verification.

        Returns:
            (tuple): Loss and dictionary of accuracies that is saved as information to the run by SMAC3.
        """
        config_hash = get_config_hash(config, 32)
        seed_ctrain(seed)

        optimizer_func = self._optimizer_from_config(config)

        model_wrapper = SABRModelWrapper(
            model=copy.deepcopy(self.original_model),
            input_shape=self.input_shape,
            eps=self.eps,
            num_epochs=epochs,
            bound_opts=self.bound_opts,
            population_bn=self.population_bn,
            relu_upper_retention=self.relu_upper_retention,
            checkpoint_save_path=None,
            device=self.device,
            train_eps_factor=config['train_eps_factor'],
            optimizer_func=optimizer_func,
            lr=config['learning_rate'],
            lr_scheduler_func=torch.optim.lr_scheduler.MultiStepLR,
            warm_up_epochs=config['warm_up_epochs'],
            ramp_up_epochs=config['ramp_up_epochs'],
            gradient_clip=self.gradient_clip,
            lr_decay_kwargs=self._lr_decay_kwargs_from_config(config, epochs),
            l1_reg_weight=config['l1_reg_weight'],
            shi_reg_weight=config['shi_reg_weight'],
            shi_reg_decay=config['shi_reg_decay'],
            sabr_subselection_ratio=config['sabr:subselection_ratio'],
            pgd_alpha=config['sabr:pgd_alpha'],
            pgd_early_stopping=False,
            pgd_restarts=config['sabr:pgd_restarts'],
            pgd_steps=config['sabr:pgd_steps'],
            pgd_decay_milestones=(),
            pgd_eps_factor=config['sabr:pgd_eps_factor'],
        )

        model_wrapper.train_model(train_loader=train_loader)
        torch.save(model_wrapper.state_dict(), f'{output_dir}/nets/{config_hash}.pt')
        model_wrapper.eval()
        std_acc, cert_acc, adv_acc = self._evaluate_hpo_model(
            model_wrapper, val_loader, cert_eval_samples, output_dir, config_hash, complete_verify
        )

        loss = 0
        loss -= nat_loss_weight * std_acc
        loss -= adv_loss_weight * adv_acc
        loss -= cert_loss_weight * cert_acc

        return loss, {'nat_acc': std_acc, 'adv_acc': adv_acc, 'cert_acc': cert_acc}

__init__(model, input_shape, eps, num_epochs, train_eps_factor=1, optimizer_func=torch.optim.Adam, lr=0.0005, warm_up_epochs=1, ramp_up_epochs=70, lr_scheduler_func=torch.optim.lr_scheduler.MultiStepLR, lr_decay_kwargs=dict(milestones=(80, 90), gamma=0.2), gradient_clip=10, l1_reg_weight=1e-06, shi_reg_weight=0.5, shi_reg_decay=True, sabr_subselection_ratio=0.2, pgd_steps=8, pgd_alpha=0.5, pgd_restarts=1, pgd_early_stopping=False, pgd_alpha_decay_factor=0.1, pgd_decay_milestones=(4, 7), pgd_eps_factor=1, checkpoint_save_path=None, checkpoint_save_interval=10, bound_opts=dict(conv_mode='patches', relu='adaptive'), device=torch.device('cuda'), population_bn=False, relu_upper_retention=1.0)

Configure the model, training schedule and SABR hyperparameters.

Parameters:

Name Type Description Default
model Module

The model to be trained.

 required
input_shape tuple

Shape of the input data.

 required
eps float

Epsilon value describing the perturbation the network should be certifiably robust against.

 required
num_epochs int

Number of epochs for training.

 required
train_eps_factor float

Factor for training epsilon.

 1
optimizer_func Optimizer

Optimizer function.

 Adam
lr float

Learning rate.

 0.0005
warm_up_epochs int

Number of warm-up epochs, i.e. epochs where the model is trained on clean loss.

 1
ramp_up_epochs int

Number of ramp-up epochs, i.e. epochs where the epsilon is gradually increased to the target train epsilon.

 70
lr_scheduler_func callable

Constructor for the learning-rate scheduler.

 MultiStepLR
lr_decay_kwargs dict

Scheduler arguments, such as milestones and gamma for MultiStepLR.

 dict(milestones=(80, 90), gamma=0.2)
gradient_clip float

Gradient clipping value.

 10
l1_reg_weight float

L1 regularisation weight.

 1e-06
shi_reg_weight float

SHI regularisation weight.

 0.5
shi_reg_decay bool

Whether to decay SHI regularisation during the ramp up phase.

 True
sabr_subselection_ratio float

Ratio of subselection for SABR (lambda in original publication).

 0.2
pgd_steps int

Number of PGD steps for adversarial loss computation.

 8
pgd_alpha float

PGD step size for adversarial loss calculation.

 0.5
pgd_restarts int

Number of PGD restarts for adversarial loss calculation.

 1
pgd_early_stopping bool

Whether to use early stopping in PGD during adversarial loss calculation.

 False
pgd_alpha_decay_factor float

PGD alpha decay factor.

 0.1
pgd_decay_milestones tuple

Milestones for PGD alpha decay.

 (4, 7)
pgd_eps_factor float

Factor for epsilon used during PGD attack.

 1
checkpoint_save_path str

Path to save checkpoints.

 None
checkpoint_save_interval int

Interval for saving checkpoints.

 10
bound_opts dict

Options for bounding according to the auto_LiRPA documentation.

 dict(conv_mode='patches', relu='adaptive')
device device

Device to run the training on.

 device('cuda')
population_bn bool

Recalibrate BatchNorm on training data after every epoch.

 False
relu_upper_retention float

Fraction of unstable ReLU upper bounds retained during training; default 1.

 1.0
Source code in CTRAIN/model_wrappers/certified/sabr_model_wrapper.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def __init__(self, model, input_shape, eps, num_epochs, train_eps_factor=1, optimizer_func=torch.optim.Adam, lr=0.0005, warm_up_epochs=1, ramp_up_epochs=70,
             lr_scheduler_func=torch.optim.lr_scheduler.MultiStepLR, lr_decay_kwargs=dict(milestones=(80, 90), gamma=0.2), gradient_clip=10, l1_reg_weight=0.000001,
             shi_reg_weight=.5, shi_reg_decay=True, sabr_subselection_ratio=.2, pgd_steps=8,
             pgd_alpha=0.5, pgd_restarts=1, pgd_early_stopping=False, pgd_alpha_decay_factor=.1,
             pgd_decay_milestones=(4,7), pgd_eps_factor=1, checkpoint_save_path=None, checkpoint_save_interval=10,
             bound_opts=dict(conv_mode='patches', relu='adaptive'), device=torch.device('cuda'), population_bn=False, relu_upper_retention=1.0):
    """Configure the model, training schedule and SABR hyperparameters.

    Args:
        model (torch.nn.Module): The model to be trained.
        input_shape (tuple): Shape of the input data.
        eps (float): Epsilon value describing the perturbation the network should be
            certifiably robust against.
        num_epochs (int): Number of epochs for training.
        train_eps_factor (float): Factor for training epsilon.
        optimizer_func (torch.optim.Optimizer): Optimizer function.
        lr (float): Learning rate.
        warm_up_epochs (int): Number of warm-up epochs, i.e. epochs where the model is
            trained on clean loss.
        ramp_up_epochs (int): Number of ramp-up epochs, i.e. epochs where the epsilon is
            gradually increased to the target train epsilon.
        lr_scheduler_func (callable): Constructor for the learning-rate scheduler.
        lr_decay_kwargs (dict): Scheduler arguments, such as milestones and gamma for MultiStepLR.
        gradient_clip (float): Gradient clipping value.
        l1_reg_weight (float): L1 regularisation weight.
        shi_reg_weight (float): SHI regularisation weight.
        shi_reg_decay (bool): Whether to decay SHI regularisation during the ramp up phase.
        sabr_subselection_ratio (float): Ratio of subselection for SABR (lambda in original
            publication).
        pgd_steps (int): Number of PGD steps for adversarial loss computation.
        pgd_alpha (float): PGD step size for adversarial loss calculation.
        pgd_restarts (int): Number of PGD restarts for adversarial loss calculation.
        pgd_early_stopping (bool): Whether to use early stopping in PGD during adversarial
            loss calculation.
        pgd_alpha_decay_factor (float): PGD alpha decay factor.
        pgd_decay_milestones (tuple): Milestones for PGD alpha decay.
        pgd_eps_factor (float): Factor for epsilon used during PGD attack.
        checkpoint_save_path (str): Path to save checkpoints.
        checkpoint_save_interval (int): Interval for saving checkpoints.
        bound_opts (dict): Options for bounding according to the auto_LiRPA documentation.
        device (torch.device): Device to run the training on.
        population_bn (bool): Recalibrate BatchNorm on training data after every epoch.
        relu_upper_retention (float): Fraction of unstable ReLU upper bounds retained during
            training; default 1.
    """
    validate_hybrid_options(relu_upper_retention, pgd_eps_factor)
    super().__init__(model, eps, input_shape, train_eps_factor, lr, optimizer_func, lr_scheduler_func, lr_decay_kwargs, bound_opts, device, checkpoint_save_path=checkpoint_save_path, checkpoint_save_interval=checkpoint_save_interval, population_bn=population_bn)
    self.relu_upper_retention = relu_upper_retention
    self.cert_train_method = 'sabr'
    self.num_epochs = num_epochs
    self.lr = lr
    self.warm_up_epochs = warm_up_epochs
    self.ramp_up_epochs = ramp_up_epochs
    self.gradient_clip = gradient_clip
    self.l1_reg_weight = l1_reg_weight
    self.shi_reg_weight = shi_reg_weight
    self.shi_reg_decay = shi_reg_decay
    self.optimizer_func = optimizer_func
    self.sabr_subselection_ratio = sabr_subselection_ratio
    self.pgd_steps = pgd_steps
    self.pgd_alpha = pgd_alpha
    self.pgd_restarts = pgd_restarts
    self.pgd_early_stopping = pgd_early_stopping
    self.pgd_alpha_decay_factor = pgd_alpha_decay_factor
    self.pgd_decay_milestones = pgd_decay_milestones
    self.pgd_eps_factor = pgd_eps_factor

train_model(train_loader, val_loader=None, start_epoch=0, end_epoch=None)

Train the wrapped model using the configured method and schedule.

The wrapper forwards its configuration to the corresponding library trainer. start_epoch and end_epoch select a segment of the full schedule, retaining its original warm-up, ramp-up and learning-rate milestones.

Parameters:

Name Type Description Default
train_loader DataLoader

DataLoader for training data.

 required
val_loader DataLoader

DataLoader for validation data.

 None
start_epoch int

Epoch to start training from. Initialises learning rate and epsilon schedulers accordingly.

 0
end_epoch int

Epoch to prematurely end training at.

 None

Returns:

Type Description
BoundedModule

Trained model.
Source code in CTRAIN/model_wrappers/certified/sabr_model_wrapper.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def train_model(self, train_loader, val_loader=None, start_epoch=0, end_epoch=None):
    """Train the wrapped model using the configured method and schedule.

    The wrapper forwards its configuration to the corresponding library trainer.
    start_epoch and end_epoch select a segment of the full schedule, retaining
    its original warm-up, ramp-up and learning-rate milestones.

    Args:
        train_loader (torch.utils.data.DataLoader): DataLoader for training data.
        val_loader (torch.utils.data.DataLoader, optional): DataLoader for validation data.
        start_epoch (int, optional): Epoch to start training from. Initialises learning rate
            and epsilon schedulers accordingly.
        end_epoch (int, optional): Epoch to prematurely end training at.

    Returns:
        (auto_LiRPA.BoundedModule): Trained model.
    """
    eps_std = self.train_eps / train_loader.std if train_loader.normalised else torch.tensor(self.train_eps)
    eps_std = torch.reshape(eps_std, (*eps_std.shape, 1, 1))
    trained_model = sabr_train_model(
        original_model=self.original_model,
        hardened_model=self.bounded_model,
        train_loader=train_loader,
        val_loader=val_loader,
        start_epoch=start_epoch,
        end_epoch=end_epoch,
        num_epochs=self.num_epochs,
        eps=self.train_eps,
        eps_std=eps_std,
        eps_schedule=(self.warm_up_epochs, self.ramp_up_epochs),
        eps_scheduler_args={},
        optimizer=self.optimizer,
        lr_scheduler=self.lr_scheduler,
        n_classes=self.n_classes,
        gradient_clip=self.gradient_clip,
        l1_regularisation_weight=self.l1_reg_weight,
        shi_regularisation_weight=self.shi_reg_weight,
        shi_reg_decay=self.shi_reg_decay,
        eps_schedule_unit='epoch',
        lr_decay_schedule_unit='epoch',
        subselection_ratio=self.sabr_subselection_ratio,
        pgd_steps=self.pgd_steps,
        pgd_step_size=self.pgd_alpha,
        pgd_restarts=self.pgd_restarts,
        pgd_early_stopping=self.pgd_early_stopping,
        pgd_decay_checkpoints=self.pgd_decay_milestones,
        pgd_decay_factor=self.pgd_alpha_decay_factor,
        pgd_eps_factor=self.pgd_eps_factor,
        results_path=self.checkpoint_path,
        checkpoint_save_interval=self.checkpoint_save_interval,
        population_bn=self.population_bn,
        relu_upper_retention=self.relu_upper_retention,
        device=self.device, selection_state=self.selection_state, history=self.history, validation_eps=self.eps)

    self.epoch = self.num_epochs if end_epoch is None else end_epoch
    return trained_model