Train CC-Dist using the shared expressive IBP training procedure.
Arguments follow expressive_ibp_train_model; this entry point selects the
method-specific loss while reusing its schedules and training loop.
Source code in CTRAIN/train/certified/cc_dist.py
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23 | def cc_dist_train_model(*args, teacher_features=None, distillation_weight=None, **kwargs):
"""Train CC-Dist using the shared expressive IBP training procedure.
Arguments follow expressive_ibp_train_model; this entry point selects the
method-specific loss while reusing its schedules and training loop.
"""
model = kwargs['hardened_model'] if 'hardened_model' in kwargs else args[1]
validate_cc_dist_bn_options(model, model.bound_opts)
if 'loss_func' not in kwargs:
if teacher_features is None:
raise ValueError('CC-Dist requires teacher_features')
kwargs['loss_func'] = partial(get_cc_dist_loss, teacher_features=teacher_features,
distillation_weight=distillation_weight)
kwargs['method'] = 'cc_dist'
kwargs['clean_forward'] = False
return expressive_ibp_train_model(*args, **kwargs)
|