def train_step(self, train_loader):
device = self.args.device
self.model.train()
CriticAcc = AverageMeter()
BaselineAcc = AverageMeter()
ActorLoss = AverageMeter()
b_loader = tqdm(train_loader)
for x_batch, y_batch, _ in b_loader:
b_loader.set_description(
f"EpochProvision: Critic: {CriticAcc.avg}, Baseline: {BaselineAcc.avg}, Actor: {ActorLoss.avg}"
)
x_batch, y_batch = x_batch.to(device), y_batch.to(device)
# Select a random batch of samples
self.optimizer.zero_grad()
labels = torch.argmax(y_batch, dim=1).long()
# Generate a batch of selections
selection_probability = self.model(x_batch, fw_module="selector")
selection = torch.bernoulli(selection_probability).detach()
# Predictor objective
critic_input = x_batch * selection
critic_out = self.model(critic_input, fw_module="predictor")
critic_loss = self.critic_loss(critic_out, labels)
# Baseline objective
baseline_out = self.model(x_batch, fw_module="baseline")
baseline_loss = self.baseline_loss(baseline_out, labels)
batch_data = torch.cat([
selection.clone().detach(),
self.softmax(critic_out).clone().detach(),
self.softmax(baseline_out).clone().detach(),
y_batch.float()
],
dim=1)
# Actor objective
actor_output = self.model(x_batch, fw_module="selector")
actor_loss = self.actor_loss(batch_data, actor_output)
total_loss = actor_loss + critic_loss + baseline_loss
total_loss.backward()
self.optimizer.step()
N = labels.shape[0]
critic_acc = accuracy(critic_out, labels)[0]
baseline_acc = accuracy(baseline_out, labels)[0]
CriticAcc.update(critic_acc.detach().item(), N)
BaselineAcc.update(baseline_acc.detach().item(), N)
ActorLoss.update(actor_loss.detach().item(), N)
summary = {
"CriticAcc": CriticAcc.avg,
"BaselineAcc": BaselineAcc.avg,
"ActorLoss": ActorLoss.avg
}
return summary
def cla_eval(self,
outs,
target,
detach=True,
criterion=nn.CrossEntropyLoss()):
'''
Args:
outs
target
criterion
Returns:
tuple: loss, prec
'''
if self.classifier is not None:
output = self.classifier(outs)
else:
output = outs
loss = criterion(output, target)
prec1 = helpers.accuracy(output, target)[0]
return loss, prec1
def custom_acc(logits, labels):
"""Returns the top1 accuracy, weighted by the class distribution.
This is important when evaluating an unbalanced dataset.
"""
batch_size = labels.size(0)
maxk = min(5, logits.shape[-1])
prec1, _ = helpers.accuracy(logits, labels, topk=(1, maxk), exact=True)
normal_prec1 = prec1.sum(0, keepdim=True).mul_(100 / batch_size)
weighted_prec1 = prec1 * class_weights[labels.cpu()].cuda()
weighted_prec1 = weighted_prec1.sum(0, keepdim=True).mul_(100 / batch_size)
return weighted_prec1.item(), normal_prec1.item()
def forward(self, x, target, *_, constraint, eps, step_size, iterations,
random_start=False, random_restarts=False, do_tqdm=False,
targeted=False, custom_loss=None, should_normalize=True,
orig_input=None, use_best=True, return_image=True, est_grad=None):
"""
Implementation of forward (finds adversarial examples). Note that
this does **not** perform inference and should not be called
directly; refer to :meth:`robustness.attacker.AttackerModel.forward`
for the function you should actually be calling.
Args:
x, target (ch.tensor) : see :meth:`robustness.attacker.AttackerModel.forward`
constraint
("2"|"inf"|"unconstrained"|"fourier"|:class:`~robustness.attack_steps.AttackerStep`)
: threat model for adversarial attacks (:math:`\ell_2` ball,
:math:`\ell_\infty` ball, :math:`[0, 1]^n`, Fourier basis, or
custom AttackerStep subclass).
eps (float) : radius for threat model.
step_size (float) : step size for adversarial attacks.
random_start (bool) : if True, start the attack with a random step.
random_restarts (bool) : if True, do many random restarts and
take the worst attack (in terms of loss) per input.
targeted (bool) : if True (False), minimize (maximize) the loss.
custom_loss (function|None) : if provided, used instead of the
criterion as the loss to maximize/minimize during
adversarial attack. The function should take in
:samp:`model, x, target` and return a tuple of the form
:samp:`loss, None`, where loss is a tensor of size N
(per-element loss).
should_normalize (bool) : If False, don't normalize the input
(not recommended unless normalization is done in the
custom_loss instead).
orig_input (ch.tensor|None) : If not None, use this as the
center of the perturbation set, rather than :samp:`x`.
return_image (bool) : If True (default), then return the adversarial
example as an image, otherwise return it in its parameterization
(for example, the Fourier coefficients if 'constraint' is
'fourier')
est_grad (tuple|None) : If not None (default), then these are
:samp:`(query_radius [R], num_queries [N])` to use for estimating the
gradient instead of autograd. We use the spherical gradient
estimator, shown below, along with antithetic sampling [#f1]_
to reduce variance:
:math:`\\nabla_x f(x) \\approx \\sum_{i=0}^N f(x + R\\cdot
\\vec{\\delta_i})\\cdot \\vec{\\delta_i}`, where
:math:`\delta_i` are randomly sampled from the unit ball.
Returns:
An adversarial example for x (i.e. within a feasible set
determined by `eps` and `constraint`, but classified as:
* `target` (if `targeted == True`)
* not `target` (if `targeted == False`)
.. [#f1] This means that we actually draw :math:`N/2` random vectors
from the unit ball, and then use :math:`\delta_{N/2+i} =
-\delta_{i}`.
"""
# Can provide a different input to make the feasible set around
# instead of the initial point
if orig_input is None: orig_input = x.detach()
orig_input = orig_input.cuda()
# Multiplier for gradient ascent [untargeted] or descent [targeted]
m = -1 if targeted else 1
# Initialize step class and attacker criterion
criterion = ch.nn.CrossEntropyLoss(reduction='none').cuda()
step_class = STEPS[constraint] if isinstance(constraint, str) else constraint
step = step_class(eps=eps, orig_input=orig_input, step_size=step_size)
def calc_loss(inp, target):
'''
Calculates the loss of an input with respect to target labels
Uses custom loss (if provided) otherwise the criterion
'''
if should_normalize:
inp = self.normalize(inp)
output = self.model(inp)
if custom_loss:
return custom_loss(self.model, inp, target)
return criterion(output, target), output
# Main function for making adversarial examples
def get_adv_examples(x):
# Random start (to escape certain types of gradient masking)
if random_start:
x = step.random_perturb(x)
# Fast Adversaril with FGSM
x = x.clone().detach()
δ = step.random_perturb(ch.zeros_like(x)).requires_grad_(True)
losses, _ = calc_loss(step.to_image(x + δ), target)
assert losses.shape[0] == x.shape[0], \
'Shape of losses must match input!'
loss = ch.mean(losses)
if step.use_grad:
if est_grad is None:
grad, = ch.autograd.grad(m * loss, [δ])
else:
f = lambda _x, _y: m * calc_loss(step.to_image(_x), _y)[0]
grad = helpers.calc_est_grad(f, xd, target, *est_grad)
else:
grad = None
with ch.no_grad():
δ = step.step(δ, grad) # δ = δ + α ∇_δ (f(x + δ))
x = step.project(x + δ)
# Save computation (don't compute last loss) if not use_best
ret = x.clone().detach()
return step.to_image(ret) if return_image else ret
# Random restarts: repeat the attack and find the worst-case
# example for each input in the batch
if random_restarts:
to_ret = None
orig_cpy = x.clone().detach()
for _ in range(random_restarts):
adv = get_adv_examples(orig_cpy)
if to_ret is None:
to_ret = adv.detach()
_, output = calc_loss(adv, target)
corr, = helpers.accuracy(output, target, topk=(1,), exact=True)
corr = corr.byte()
misclass = ~corr
to_ret[misclass] = adv[misclass]
adv_ret = to_ret
else:
adv_ret = get_adv_examples(x)
return adv_ret
def _model_loop(args, loop_type, loader, model, opt, epoch, adv, writer):
if not loop_type in ['train', 'val']:
err_msg = "loop_type ({0}) must be 'train' or 'val'".format(loop_type)
raise ValueError(err_msg)
is_train = (loop_type == 'train')
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
prec = 'NatPrec' if not adv else 'AdvPrec'
loop_msg = 'Train' if loop_type == 'train' else 'Val'
# switch to train/eval mode depending
model = model.train() if is_train else model.eval()
# If adv training (or evaling), set eps and random_restarts appropriately
if adv:
eps = calc_fadein_eps(epoch, args.eps_fadein_epochs, args.eps) \
if is_train else args.eps
random_restarts = 0 if is_train else args.random_restarts
# Custom training criterion
has_custom_train_loss = has_attr(args, 'custom_train_loss')
train_criterion = args.custom_train_loss if has_custom_train_loss \
else ch.nn.CrossEntropyLoss()
has_custom_adv_loss = has_attr(args, 'custom_adv_loss')
adv_criterion = args.custom_adv_loss if has_custom_adv_loss else None
attack_kwargs = {}
if adv:
if loop_type == 'train':
attack_kwargs = {
'constraint': args.constraint,
'eps': eps,
'iterations': 1,
'step_size': args.fast_attack_lr,
'random_start': args.random_start,
'custom_loss': adv_criterion,
'random_restarts': random_restarts,
}
elif loop_type == 'val':
attack_kwargs = {
'constraint': args.constraint,
'eps': eps,
'step_size': args.attack_lr,
'iterations': args.attack_steps,
'random_start': args.random_start,
'custom_loss': adv_criterion,
'random_restarts': random_restarts,
'use_best': bool(args.use_best)
}
iterator = tqdm(enumerate(loader), total=len(loader))
for i, (inp, target) in iterator:
# measure data loading time
target = target.cuda(non_blocking=True)
output, final_inp = model(inp,
target=target,
make_adv=adv,
**attack_kwargs)
loss = train_criterion(output, target)
if len(loss.shape) > 0: loss = loss.mean()
model_logits = output[0] if (type(output) is tuple) else output
# measure accuracy and record loss
top1_acc = float('nan')
top5_acc = float('nan')
try:
maxk = min(5, model_logits.shape[-1])
prec1, prec5 = helpers.accuracy(model_logits, target, topk=(1, maxk))
losses.update(loss.item(), inp.size(0))
top1.update(prec1[0], inp.size(0))
top5.update(prec5[0], inp.size(0))
top1_acc = top1.avg
top5_acc = top5.avg
except:
pass
reg_term = 0.0
if has_attr(args, "regularizer"):
reg_term = args.regularizer(model, inp, target)
loss = loss + reg_term
# compute gradient and do SGD step
if is_train:
opt.zero_grad()
loss.backward()
opt.step()
elif adv and i == 0 and writer:
# add some examples to the tensorboard
nat_grid = make_grid(inp[:15, ...])
adv_grid = make_grid(final_inp[:15, ...])
writer.add_image('Nat input', nat_grid, epoch)
writer.add_image('Adv input', adv_grid, epoch)
# ITERATOR
desc = ('{2} Epoch:{0} | Loss {loss.avg:.4f} | '
'{1}1 {top1_acc:.3f} | {1}5 {top5_acc:.3f} | '
'Reg term: {reg} ||'.format( epoch, prec, loop_msg,
loss=losses, top1_acc=top1_acc, top5_acc=top5_acc, reg=reg_term))
# USER-DEFINED HOOK
if has_attr(args, 'iteration_hook'):
args.iteration_hook(model, i, loop_type, inp, target)
iterator.set_description(desc)
iterator.refresh()
if writer is not None:
prec_type = 'adv' if adv else 'nat'
descs = ['loss', 'top1', 'top5']
vals = [losses, top1, top5]
for d, v in zip(descs, vals):
writer.add_scalar('_'.join([prec_type, loop_type, d]), v.avg,
epoch)
return top1.avg, losses.avg