def train_one_epoch(net,
batch_generator,
optimizer,
criterion,
device,
descr_str='Training',
attack_method=None,
adv_coef=1.0):
"""
:return: clean_acc, adv_acc
"""
net.train()
pbar = tqdm(batch_generator)
advacc = -1
advloss = -1
cleanacc = -1
cleanloss = -1
pbar.set_description(descr_str)
for i, (data, label) in enumerate(pbar):
data, label = data.to(device), label.to(device)
optimizer.zero_grad()
pbar_dic = OrderedDict()
# TotalLoss = 0
if attack_method is not None:
adv_inp = attack_method.attack(net, data, label)
optimizer.zero_grad()
net.train()
pred = net(adv_inp)
loss = criterion(pred, label)
acc = torch_accuracy(pred, label, (1, ))
advacc = acc[0].item()
advloss = loss.item()
# TotalLoss = TotalLoss + loss * adv_coef
(loss * adv_coef).backward()
pred = net(data)
loss = criterion(pred, label)
# TotalLoss = TotalLoss + loss
loss.backward()
# TotalLoss.backward()
# param = next(net.parameters())
# grad_mean = torch.mean(param.grad)
optimizer.step()
acc = torch_accuracy(pred, label, (1, ))
cleanacc = acc[0].item()
cleanloss = loss.item()
# pbar_dic['grad'] = '{}'.format(grad_mean)
pbar_dic['Acc'] = '{:.2f}'.format(cleanacc)
pbar_dic['loss'] = '{:.2f}'.format(cleanloss)
pbar_dic['AdvAcc'] = '{:.2f}'.format(advacc)
pbar_dic['Advloss'] = '{:.2f}'.format(advloss)
pbar.set_postfix(pbar_dic)
return cleanacc, advacc
def eval_one_epoch(net, batch_generator, device, attack_method=None):
net.eval()
pbar = tqdm(batch_generator)
clean_accuracy = AvgMeter()
adv_accuracy = AvgMeter()
pbar.set_description('Evaluating')
for (data, label) in pbar:
data, label = data.to(device), label.to(device)
with torch.no_grad():
h1, h2, h3, h4, y = net(data)
acc = torch_accuracy(y, label, (1, ))
clean_accuracy.update(acc[0].item())
if attack_method is not None:
adv_inp = attack_method.attack(net, data, label)
with torch.no_grad():
h1a, h2a, h3a, h4a, ya = net(adv_inp)
acc = torch_accuracy(ya, label, (1, ))
adv_accuracy.update(acc[0].item())
pbar_dic = OrderedDict()
pbar_dic['CleanAcc'] = '{:.2f}'.format(clean_accuracy.mean)
pbar_dic['AdvAcc'] = '{:.2f}'.format(adv_accuracy.mean)
pbar.set_postfix(pbar_dic)
adv_acc = adv_accuracy.mean if attack_method is not None else 0
return clean_accuracy.mean, adv_acc
def train_hloss(net,
batch_generator,
optimizer,
criterion,
device,
descr_str='Training',
attack_method=None,
adv_coef=1.0):
"""
:return: clean_acc, adv_acc
"""
net.train()
pbar = tqdm(batch_generator)
advacc = -1
cleanacc = -1
mse_criterion = nn.MSELoss()
pbar.set_description(descr_str)
for i, (data, label) in enumerate(pbar):
data, label = data.to(device), label.to(device)
optimizer.zero_grad()
pbar_dic = OrderedDict()
if attack_method is not None:
adv_inp = attack_method.attack(net, data, label)
optimizer.zero_grad()
net.train()
h1a, h2a, h3a, h4a, ya = net(adv_inp)
xent_loss = criterion(ya, label)
acc = torch_accuracy(ya, label, (1, ))
advacc = acc[0].item()
h1, h2, h3, h4, y = net(data)
h_loss = mse_criterion(h4a, h4)
# h_loss = kl_div_loss(h4a, h4.detach(), 1)
loss = h_loss * adv_coef + xent_loss
loss.sum().backward()
optimizer.step()
acc = torch_accuracy(y, label, (1, ))
cleanacc = acc[0].item()
pbar_dic['Acc'] = '{:.2f}'.format(cleanacc)
pbar_dic['hloss'] = '{:.2f}'.format(h_loss.item())
pbar_dic['AdvAcc'] = '{:.2f}'.format(advacc)
pbar_dic['xentloss'] = '{:.2f}'.format(xent_loss.item())
pbar.set_postfix(pbar_dic)
return cleanacc, advacc
def train_mi(net,
net2,
batch_generator,
optimizer,
criterion,
device,
descr_str='Training',
attack_method=None,
adv_coef=1.0):
"""
:return: clean_acc, adv_acc
"""
net.train()
pbar = tqdm(batch_generator)
advacc = -1
cleanacc = -1
pbar.set_description(descr_str)
for i, (data, label) in enumerate(pbar):
data, label = data.to(device), label.to(device)
optimizer.zero_grad()
pbar_dic = OrderedDict()
if attack_method is not None:
adv_inp = attack_method.attack(net, data, label)
optimizer.zero_grad()
net.train()
h1a, h2a, h3a, h4a, ya = net(adv_inp)
xent_loss = criterion(ya, label)
h1_prime = torch.cat((h1a[1:], h1a[0].unsqueeze(0)), dim=0)
mi_loss = net2(h4a, h1a, h1_prime)
acc = torch_accuracy(ya, label, (1, ))
advacc = acc[0].item()
h1, h2, h3, h4, y = net(data)
loss = mi_loss.sum() * adv_coef + xent_loss.sum()
loss.sum().backward()
optimizer.step()
acc = torch_accuracy(y, label, (1, ))
cleanacc = acc[0].item()
pbar_dic['Acc'] = '{:.2f}'.format(cleanacc)
pbar_dic['miloss'] = '{:.2f}'.format(mi_loss.sum().item())
pbar_dic['AdvAcc'] = '{:.2f}'.format(advacc)
pbar_dic['xentloss'] = '{:.2f}'.format(xent_loss.sum().item())
pbar.set_postfix(pbar_dic)
return cleanacc, advacc
def eval_one_epoch(net,
batch_generator,
DEVICE=torch.device('cuda:0'),
AttackMethod=None):
net.eval()
pbar = tqdm(batch_generator)
clean_accuracy = AvgMeter()
adv_accuracy = AvgMeter()
pbar.set_description('Evaluating')
for (data, label) in pbar:
# data = data.to(DEVICE)
# label = label.to(DEVICE)
data = data.cuda()
label = label.cuda()
# with torch.no_grad():
pred = net(data)
acc = torch_accuracy(pred, label, (1, ))
clean_accuracy.update(acc[0].item())
if AttackMethod is not None:
adv_inp = AttackMethod.attack(net, data, label)
# with torch.no_grad():
pred = net(adv_inp)
acc = torch_accuracy(pred, label, (1, ))
adv_accuracy.update(acc[0].item())
pbar_dic = OrderedDict()
pbar_dic['CleanAcc'] = '{:.2f}'.format(clean_accuracy.mean)
pbar_dic['AdvAcc'] = '{:.2f}'.format(adv_accuracy.mean)
pbar.set_postfix(pbar_dic)
adv_acc = adv_accuracy.mean if AttackMethod is not None else 0
return clean_accuracy.mean, adv_acc
def train_one_epoch(net,
batch_generator,
optimizer,
criterion,
DEVICE=torch.device('cuda:0'),
descrip_str='Training',
AttackMethod=None,
oracle=None,
adv_coef=1.0):
'''
:param attack_freq: Frequencies of training with adversarial examples. -1 indicates natural training
:param AttackMethod: the attack method, None represents natural training
:return: None #(clean_acc, adv_acc)
'''
net.train()
pbar = tqdm(batch_generator)
advacc = -1
advloss = -1
cleanacc = -1
cleanloss = -1
pbar.set_description(descrip_str)
for i, (data, label) in enumerate(pbar):
# data = data.to(DEVICE)
data = data.cuda()
label = label.cuda()
# label = label.to(DEVICE)
optimizer.zero_grad()
pbar_dic = OrderedDict()
TotalLoss = 0
if AttackMethod is not None:
if oracle is None:
adv_inp = AttackMethod.attack(net, data, label)
else:
adv_inp = AttackMethod.attack(oracle, data, label)
optimizer.zero_grad()
net.train()
pred = net(adv_inp)
loss = criterion(pred, label)
acc = torch_accuracy(pred, label, (1, ))
advacc = acc[0].item()
advloss = loss.item()
#TotalLoss = TotalLoss + loss * adv_coef
(loss * adv_coef).backward()
pred = net(data)
loss = criterion(pred, label)
#TotalLoss = TotalLoss + loss
loss.backward()
#TotalLoss.backward()
#param = next(net.parameters())
#grad_mean = torch.mean(param.grad)
optimizer.step()
acc = torch_accuracy(pred, label, (1, ))
cleanacc = acc[0].item()
cleanloss = loss.item()
#pbar_dic['grad'] = '{}'.format(grad_mean)
pbar_dic['Acc'] = '{:.2f}'.format(cleanacc)
pbar_dic['loss'] = '{:.2f}'.format(cleanloss)
pbar_dic['AdvAcc'] = '{:.2f}'.format(advacc)
pbar_dic['Advloss'] = '{:.2f}'.format(advloss)
pbar.set_postfix(pbar_dic)
def train_one_epoch(net, batch_generator, optimizer,
criterion, LayerOneTrainner, K,
device='cuda:0', descr_str='Training'):
"""
:param attack_freq: Frequencies of training with adversarial examples. -1 indicates natural training
:param AttackMethod: the attack method, None represents natural training
:return: None #(clean_acc, adv_acc)
"""
net.train()
pbar = tqdm(batch_generator)
yofoacc = -1
cleanacc = -1
cleanloss = -1
pbar.set_description(descr_str)
for i, (data, label) in enumerate(pbar):
data = data.to(device)
label = label.to(device)
eta = torch.FloatTensor(*data.shape).uniform_(-config.eps, config.eps)
eta = eta.to(label.device)
eta.requires_grad_()
optimizer.zero_grad()
LayerOneTrainner.param_optimizer.zero_grad()
for j in range(K):
# optimizer.zero_grad()
pbar_dic = OrderedDict()
TotalLoss = 0
pred = net(data + eta.detach())
loss = criterion(pred, label)
TotalLoss = TotalLoss + loss
# wgrad = net.conv1.weight.grad
# bgrad = net.conv1.bias.grad
TotalLoss.backward()
# net.conv1.weight.grad = wgrad
# net.conv1.bias.grad = bgrad
# param = next(net.parameters())
# grad_mean = torch.mean(param.grad)
# optimizer.step()
# optimizer.zero_grad()
p = -1.0 * net.layer_one_out.grad
yofo_inp, eta = LayerOneTrainner.step(data, p, eta)
with torch.no_grad():
if j == 0:
acc = torch_accuracy(pred, label, (1,))
cleanacc = acc[0].item()
cleanloss = loss.item()
if j == K - 1:
yofo_pred = net(yofo_inp)
yofoacc = torch_accuracy(yofo_pred, label, (1,))[0].item()
# pbar_dic['grad'] = '{}'.format(grad_mean)
optimizer.step()
LayerOneTrainner.param_optimizer.step()
optimizer.zero_grad()
LayerOneTrainner.param_optimizer.zero_grad()
pbar_dic['Acc'] = '{:.2f}'.format(cleanacc)
pbar_dic['loss'] = '{:.2f}'.format(cleanloss)
pbar_dic['YofoAcc'] = '{:.2f}'.format(yofoacc)
pbar.set_postfix(pbar_dic)
return cleanacc, yofoacc
def train_one_epoch(net,
batch_generator,
optimizer,
criterion,
LayerOneTrainer,
K,
device,
descr_str='Training'):
"""
:return: clean_acc, adv_acc
"""
net.train()
pbar = tqdm(batch_generator)
yofoacc = -1
cleanacc = -1
cleanloss = -1
pbar.set_description(descr_str)
for i, (data, label) in enumerate(pbar):
data, label = data.to(device), label.to(device)
eta = torch.cuda.FloatTensor(*data.shape).uniform_(
-config.eps, config.eps)
eta = eta.to(device)
eta.requires_grad_()
optimizer.zero_grad()
LayerOneTrainer.param_optimizer.zero_grad()
for j in range(K):
# optimizer.zero_grad()
pbar_dic = OrderedDict()
TotalLoss = 0
h1, h2, h3, h4, y = net(data + eta.detach())
loss = criterion(y, label)
TotalLoss = TotalLoss + loss
wgrad = net.conv1.weight.grad
# bgrad = net.conv1.bias.grad
TotalLoss.backward()
net.conv1.weight.grad = wgrad
# net.conv1.bias.grad = bgrad
# param = next(net.parameters())
# grad_mean = torch.mean(param.grad)
# optimizer.step()
# optimizer.zero_grad()
p = -1.0 * net.layer_one_out.grad
yofo_inp, eta = LayerOneTrainer.step(data, p, eta)
with torch.no_grad():
if j == 0:
acc = torch_accuracy(y, label, (1, ))
cleanacc = acc[0].item()
cleanloss = loss.item()
if j == K - 1:
h1, h2, h3, h4, y = net(yofo_inp)
yofoacc = torch_accuracy(y, label, (1, ))[0].item()
# pbar_dic['grad'] = '{}'.format(grad_mean)
optimizer.step()
LayerOneTrainer.param_optimizer.step()
optimizer.zero_grad()
LayerOneTrainer.param_optimizer.zero_grad()
pbar_dic['Acc'] = '{:.2f}'.format(cleanacc)
pbar_dic['loss'] = '{:.2f}'.format(cleanloss)
pbar_dic['YofoAcc'] = '{:.2f}'.format(yofoacc)
pbar.set_postfix(pbar_dic)
return cleanacc, yofoacc
def train_one_epoch(net, batch_generator, optimizer,
criterion, LayerOneTrainner, K,
device='cuda', descr_str='Training'):
net.train()
pbar = tqdm(batch_generator)
yofoacc = -1
pbar.set_description(descr_str)
trades_criterion = torch.nn.KLDivLoss(reduction='sum') # .to(DEVICE)
for i, (data, label) in enumerate(pbar):
data = data.to(device)
label = label.to(device)
net.eval()
eta = 0.001 * torch.randn(data.shape).cuda().detach().to(device)
eta.requires_grad_()
raw_soft_label = F.softmax(net(data), dim=1).detach()
for j in range(K):
pred = net(data + eta.detach())
with torch.enable_grad():
loss = trades_criterion(F.log_softmax(pred, dim=1), raw_soft_label) # raw_soft_label.detach())
p = -1.0 * torch.autograd.grad(loss, [net.layer_one_out, ])[0]
yofo_inp, eta = LayerOneTrainner.step(data, p, eta)
with torch.no_grad():
if j == K - 1:
yofo_pred = net(yofo_inp)
yofo_loss = criterion(yofo_pred, label)
yofoacc = torch_accuracy(yofo_pred, label, (1,))[0].item()
net.train()
optimizer.zero_grad()
LayerOneTrainner.param_optimizer.zero_grad()
raw_pred = net(data)
acc = torch_accuracy(raw_pred, label, (1,))
clean_acc = acc[0].item()
clean_loss = criterion(raw_pred, label)
adv_pred = net(torch.clamp(data + eta.detach(), 0.0, 1.0))
kl_loss = trades_criterion(F.log_softmax(adv_pred, dim=1),
F.softmax(raw_pred, dim=1)) / data.shape[0]
loss = clean_loss + kl_loss
loss.backward()
optimizer.step()
LayerOneTrainner.param_optimizer.step()
optimizer.zero_grad()
LayerOneTrainner.param_optimizer.zero_grad()
pbar_dic = OrderedDict()
pbar_dic['Acc'] = '{:.2f}'.format(clean_acc)
pbar_dic['cleanloss'] = '{:.3f}'.format(clean_loss.item())
pbar_dic['klloss'] = '{:.3f}'.format(kl_loss.item())
pbar_dic['YofoAcc'] = '{:.2f}'.format(yofoacc)
pbar_dic['Yofoloss'] = '{:.3f}'.format(yofo_loss.item())
pbar.set_postfix(pbar_dic)
return clean_acc, yofoacc
def trades_loss(model,
x_natural,
y,
optimizer,
device,
step_size=0.003,
epsilon=0.031,
perturb_steps=10,
beta=1.0,
distance='l_inf'):
# define KL-loss
criterion_kl = nn.KLDivLoss(reduction='sum')
model.eval()
batch_size = len(x_natural)
# generate adversarial example
x_adv = x_natural.detach() + 0.001 * torch.randn(
x_natural.shape).cuda().detach().to(device)
if distance == 'l_inf':
# logits_natural = model(x_natural).detach()
for _ in range(perturb_steps):
x_adv.requires_grad_()
with torch.enable_grad():
loss_kl = criterion_kl(F.log_softmax(model(x_adv), dim=1),
F.softmax(model(x_natural), dim=1))
# loss_kl = criterion_kl(F.log_softmax(model(x_adv), dim=1),
# F.softmax(logits_natural, dim=1))
grad = torch.autograd.grad(loss_kl, [x_adv])[0]
x_adv = x_adv.detach() + step_size * torch.sign(grad.detach())
x_adv = torch.min(torch.max(x_adv, x_natural - epsilon),
x_natural + epsilon)
x_adv = torch.clamp(x_adv, 0.0, 1.0)
elif distance == 'l_2':
for _ in range(perturb_steps):
x_adv.requires_grad_()
with torch.enable_grad():
loss_kl = criterion_kl(F.log_softmax(model(x_adv), dim=1),
F.softmax(model(x_natural), dim=1))
grad = torch.autograd.grad(loss_kl, [x_adv])[0]
for idx_batch in range(batch_size):
grad_idx = grad[idx_batch]
grad_idx_norm = l2_norm(grad_idx)
grad_idx /= (grad_idx_norm + 1e-8)
x_adv[idx_batch] = x_adv[idx_batch].detach(
) + step_size * grad_idx
eta_x_adv = x_adv[idx_batch] - x_natural[idx_batch]
norm_eta = l2_norm(eta_x_adv)
if norm_eta > epsilon:
eta_x_adv = eta_x_adv * epsilon / l2_norm(eta_x_adv)
x_adv[idx_batch] = x_natural[idx_batch] + eta_x_adv
x_adv = torch.clamp(x_adv, 0.0, 1.0)
else:
x_adv = torch.clamp(x_adv, 0.0, 1.0)
model.train()
x_adv = Variable(torch.clamp(x_adv, 0.0, 1.0), requires_grad=False)
# zero gradient
optimizer.zero_grad()
# calculate robust loss
logits = model(x_natural)
adv_logits = model(x_adv)
loss_natural = F.cross_entropy(logits, y)
loss_robust = (1.0 / batch_size) * criterion_kl(
F.log_softmax(adv_logits, dim=1), F.softmax(logits, dim=1))
loss = loss_natural + beta * loss_robust
cleanacc = torch_accuracy(logits, y, (1, ))[0].item()
tradesacc = torch_accuracy(adv_logits, y, (1, ))[0].item()
return loss, loss_natural.item(), loss_robust.item(), cleanacc, tradesacc
