Skip to content

util

Shared validation, model selection, BatchNorm synchronisation and epoch-metric utilities.

EpochMetrics

Accumulate sample-weighted training metrics across full and partial batches.

Source code in CTRAIN/train/certified/util.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
class EpochMetrics:
    """Accumulate sample-weighted training metrics across full and partial batches."""
    def __init__(self):
        self.samples = 0
        self.totals = {}

    def update(self, count, **values):
        """Accumulate batch-mean metrics weighted by the number of examples."""
        self.samples += count
        for key, value in values.items():
            self.totals[key] = self.totals.get(key, 0.) + float(value) * count

    def mean(self, key):
        """Return the sample-weighted mean of a previously accumulated metric."""
        return self.totals[key] / self.samples

mean(key)

Return the sample-weighted mean of a previously accumulated metric.

Source code in CTRAIN/train/certified/util.py 
52
53
54
def mean(self, key):
    """Return the sample-weighted mean of a previously accumulated metric."""
    return self.totals[key] / self.samples

update(count, **values)

Accumulate batch-mean metrics weighted by the number of examples.

Source code in CTRAIN/train/certified/util.py 
46
47
48
49
50
def update(self, count, **values):
    """Accumulate batch-mean metrics weighted by the number of examples."""
    self.samples += count
    for key, value in values.items():
        self.totals[key] = self.totals.get(key, 0.) + float(value) * count

checkpoint_due(epoch, end_epoch, interval)

Return whether a completed epoch is a periodic or final trajectory checkpoint.

Source code in CTRAIN/train/certified/util.py 
35
36
37
def checkpoint_due(epoch, end_epoch, interval):
    """Return whether a completed epoch is a periodic or final trajectory checkpoint."""
    return epoch % interval == 0 or epoch == end_epoch

evaluate_taps_proxy(original_model, hardened_model, data_loader, eps, n_classes, device, propagation='IBP', sabr_args=None, **attack_args)

Evaluate latent-margin proxy accuracy at a fixed perturbation radius.

The feature and classifier split uses IBP or SABR propagation and latent PGD. Caller modes and gradient flags are preserved, so BatchNorm statistics are not updated. This score is used for model selection, not sound certification.

Source code in CTRAIN/train/certified/util.py 
 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
def evaluate_taps_proxy(original_model, hardened_model, data_loader, eps, n_classes,
                        device, propagation="IBP", sabr_args=None, **attack_args):
    """Evaluate latent-margin proxy accuracy at a fixed perturbation radius.

    The feature and classifier split uses IBP or SABR propagation and latent PGD.
    Caller modes and gradient flags are preserved, so BatchNorm statistics are
    not updated. This score is used for model selection, not sound certification.
    """
    from auto_LiRPA import PerturbationLpNorm
    from CTRAIN.bound import bound_taps
    from CTRAIN.util.util import preserve_model_state

    models = (original_model, hardened_model, *hardened_model.bounded_blocks)
    correct = total = 0
    with preserve_model_state(*models), torch.no_grad():
        for model in models:
            model.eval()
        radius = torch.as_tensor(eps, device=device).reshape(-1, 1, 1)
        data_min, data_max = data_loader.min.to(device), data_loader.max.to(device)
        for data, target in data_loader:
            data, target = data.to(device), target.to(device)
            if (radius == 0).all():
                correct += (hardened_model(data).argmax(1) == target).sum().item()
            else:
                ptb = PerturbationLpNorm(
                    norm=np.inf, eps=radius,
                    x_L=torch.clamp(data - radius, data_min, data_max),
                    x_U=torch.clamp(data + radius, data_min, data_max),
                )
                region_args = None if sabr_args is None else dict(
                    sabr_args, hardened_model=hardened_model, original_model=original_model,
                    data=data, target=target, eps=radius, data_min=data_min,
                    data_max=data_max, device=device, n_classes=n_classes,
                )
                margins, _ = bound_taps(
                    original_model, hardened_model, hardened_model.bounded_blocks,
                    data, target, n_classes, ptb, device=device,
                    propagation=propagation, sabr_args=region_args, **attack_args,
                )
                correct += (margins > 0).all(1).sum().item()
            total += len(data)
    if not total:
        raise ValueError("Validation loader must contain at least one sample")
    return correct / total

evaluate_validation_loss(models, loader, loss_func, device)

Compute sample-weighted validation loss with fixed evaluation-mode BatchNorm statistics.

Source code in CTRAIN/train/certified/util.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
def evaluate_validation_loss(models, loader, loss_func, device):
    """Compute sample-weighted validation loss with fixed evaluation-mode BatchNorm statistics."""
    from CTRAIN.util.util import preserve_model_state
    total = samples = 0
    with preserve_model_state(*models), torch.no_grad():
        for model in models:
            model.eval()
        for data, target in loader:
            data, target = data.to(device), target.to(device)
            loss = loss_func(data, target)
            total += float(loss) * len(data)
            samples += len(data)
    if not samples:
        raise ValueError('Validation loader must contain at least one sample')
    return total / samples

prepare_model_selection(state, loader, eps=None)

Prepare persistent validation-selection state and reject an empty validation loader.

Source code in CTRAIN/train/certified/util.py 
141
142
143
144
145
146
147
148
149
def prepare_model_selection(state, loader, eps=None):
    """Prepare persistent validation-selection state and reject an empty validation loader."""
    if state is not None and not isinstance(state, dict):
        raise ValueError('selection_state must be a dictionary')
    if loader is not None and not len(loader):
        raise ValueError('Validation loader must contain at least one sample')
    if eps is not None and (not math.isfinite(eps) or eps < 0):
        raise ValueError('Validation radius must be nonnegative and finite')
    return {} if state is None else state

resolve_validation_eps(validation_eps, eval_eps)

Resolve validation_eps and its legacy eval_eps alias, rejecting conflicting radii.

Source code in CTRAIN/train/certified/util.py 
28
29
30
31
32
def resolve_validation_eps(validation_eps, eval_eps):
    """Resolve validation_eps and its legacy eval_eps alias, rejecting conflicting radii."""
    if validation_eps is not None and eval_eps is not None and validation_eps != eval_eps:
        raise ValueError('validation_eps and legacy eval_eps must agree')
    return eval_eps if validation_eps is None else validation_eps

restore_model_selection(state, model)

Load selected inference weights if present and return whether restoration occurred.

Source code in CTRAIN/train/certified/util.py 
163
164
165
166
167
168
def restore_model_selection(state, model):
    """Load selected inference weights if present and return whether restoration occurred."""
    if state.get('weights') is None:
        return False
    model.load_state_dict(state['weights'])
    return True

synchronise_bn(original_model, hardened_model, train_loader, device, population=False, loss_fusion_model=None)

Synchronise BatchNorm buffers across the original, bounded and fused models.

Optional population recalibration refreshes running statistics on training data before copying them. Recalibration preserves caller modes and BatchNorm momentum, including on failure.

Source code in CTRAIN/train/certified/util.py 
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
def synchronise_bn(original_model, hardened_model, train_loader, device, population=False, loss_fusion_model=None):
    """Synchronise BatchNorm buffers across the original, bounded and fused models.

    Optional population recalibration refreshes running statistics on training
    data before copying them. Recalibration preserves caller modes and BatchNorm
    momentum, including on failure.
    """
    import torch
    from torch.nn.modules.batchnorm import _BatchNorm
    from CTRAIN.util.util import preserve_model_state

    original = original_model
    with torch.no_grad():
        buffers = dict(original.named_buffers())
        for name, value in hardened_model.named_buffers():
            buffers[name].copy_(value)
        layers = [m for m in original.modules() if isinstance(m, _BatchNorm)]
        if population and layers:
            momenta = [m.momentum for m in layers]
            try:
                with preserve_model_state(original):
                    original.train()
                    for batch, (data, _) in enumerate(train_loader, 1):
                        for layer in layers:
                            layer.momentum = 1 / batch
                        original(data.to(device))
            finally:
                for layer, momentum in zip(layers, momenta):
                    layer.momentum = momentum
        buffers = dict(original.named_buffers())
        for name, value in hardened_model.named_buffers():
            value.copy_(buffers[name])
        if loss_fusion_model is not None:
            buffers = dict(original.named_buffers())
            for name, value in loss_fusion_model.named_buffers():
                value.copy_(buffers[name.removeprefix("model.")])

update_model_selection(state, model, score, epoch, *, allow_negative_infinity=False)

Retain the first best-scoring inference weights without rewinding the training trajectory.

Source code in CTRAIN/train/certified/util.py 
152
153
154
155
156
157
158
159
160
def update_model_selection(state, model, score, epoch, *, allow_negative_infinity=False):
    """Retain the first best-scoring inference weights without rewinding the training trajectory."""
    score = tuple(float(v) for v in score) if isinstance(score, tuple) else float(score)
    values = score if isinstance(score, tuple) else (score,)
    if any(not math.isfinite(v) and not (allow_negative_infinity and v == -float('inf')) for v in values):
        raise FloatingPointError('Validation score must be finite')
    previous = state.get('score', state.get('accuracy'))
    if previous is None or score > previous:
        state.update(score=score, weights=copy.deepcopy(model.state_dict()), epoch=epoch)

validate_training_run(train_loader, num_epochs, start_epoch, end_epoch, checkpoint_save_interval, lr_decay_schedule_unit='epoch', *, history=None)

Validate the loader, epoch range, scheduler unit and checkpoint interval; return the exclusive ending epoch.

Source code in CTRAIN/train/certified/util.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
def validate_training_run(train_loader, num_epochs, start_epoch, end_epoch,
                          checkpoint_save_interval, lr_decay_schedule_unit='epoch', *, history=None):
    """Validate the loader, epoch range, scheduler unit and checkpoint interval; return the
    exclusive ending epoch.
    """
    end_epoch = num_epochs if end_epoch is None else end_epoch
    if (any(type(v) is not int for v in (num_epochs, start_epoch, end_epoch))
            or num_epochs <= 0 or not 0 <= start_epoch <= end_epoch <= num_epochs):
        raise ValueError('Require 0 <= start_epoch <= end_epoch <= num_epochs, with positive num_epochs')
    if type(checkpoint_save_interval) is not int or checkpoint_save_interval <= 0:
        raise ValueError('checkpoint_save_interval must be a positive integer')
    if not len(train_loader):
        raise ValueError('Training loader must contain batches')
    if lr_decay_schedule_unit not in ('epoch', 'batch'):
        raise ValueError("lr_decay_schedule_unit must be 'epoch' or 'batch'")
    if history is not None and not isinstance(history, list):
        raise ValueError('history must be a list')
    return end_epoch

validation_ptb(data, loader, eps)

Construct a clipped validation box from a raw-input radius and the loader's metadata.

Source code in CTRAIN/train/certified/util.py 
188
189
190
191
192
193
194
195
196
def validation_ptb(data, loader, eps):
    """Construct a clipped validation box from a raw-input radius and the loader's metadata."""
    from auto_LiRPA import PerturbationLpNorm
    from CTRAIN.util.util import loader_radius, loader_limits
    radius = loader_radius(eps, loader, data)
    minimum, maximum = loader_limits(loader, data)
    return PerturbationLpNorm(norm=float('inf'), eps=radius,
                             x_L=(data - radius).clamp(minimum, maximum),
                             x_U=(data + radius).clamp(minimum, maximum))