def train(args, model, device, train_loader, optimizer, epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
# calculate robust loss
loss = trades_loss(model=model,
x_natural=data,
y=target,
optimizer=optimizer,
step_size=args.step_size,
epsilon=args.epsilon,
perturb_steps=args.num_steps,
beta=args.beta)
loss.backward()
optimizer.step()
# print progress
if batch_idx % args.log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
def train(args, model, device, train_loader, optimizer, epoch):
model.train()
loss_his = []
loss_robust_his = []
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
# calculate robust loss
loss, loss_robust = trades_loss(model=model,
x_natural=data,
y=target,
optimizer=optimizer,
args=args)
loss.backward()
optimizer.step()
loss_his.append(loss.item())
loss_robust_his.append(loss_robust.item())
# print progress
if batch_idx % args.log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
break
return np.mean(loss_his), np.mean(loss_robust_his)
def train_step(self, batch, batch_idx):
x, y = batch
x = x.to(self.device)
y = y.to(self.device)
if self.model.training:
loss, logits, adv_logits = trades_loss(
self.model, x, y, self.optimizer, self.args.eps_iter,
self.args.eps, self.args.nb_iter, self.args.adv_loss_wt)
else:
xadv = self.compute_adversarial(x, y)
logits = self.model(x)
adv_logits = self.model(xadv)
loss = torch.nn.functional.cross_entropy(logits, y)
loss += torch.nn.functional.cross_entropy(adv_logits, y)
cln_acc, _ = compute_accuracy(logits, y)
adv_acc, _ = compute_accuracy(adv_logits, y)
if self.model.training:
for p in self.model.parameters():
assert (p.grad is None) or (p.grad == 0).all()
return {
'loss': loss
}, {
'train_clean_accuracy': cln_acc,
'train_adv_accuracy': adv_acc,
'train_accuracy': (cln_acc + adv_acc) / 2,
'train_loss': float(loss.detach().cpu()),
}
def train(args, model, device, train_loader, optimizer, epoch):
cudnn.benchmark = True
model.train()
batch_time = utils.AverageMeter()
losses = utils.AverageMeter()
end = time.time()
for batch_idx, (data1, target1, data2, target2) in enumerate(train_loader):
if batch_idx % 2 == 0:
data, target = data1.to(device), target1.to(device)
else:
data, target = data2.to(device), target2.to(device)
optimizer.zero_grad()
if args.trades:
# calculate robust loss
loss = trades_loss(model=model,
x_natural=data,
y=target,
optimizer=optimizer,
step_size=args.step_size,
epsilon=args.epsilon,
perturb_steps=args.num_steps,
beta=args.beta,
distance=args.distance,
device=device)
else:
# CE loss
loss = F.cross_entropy(model(data), target)
loss.backward()
optimizer.step()
batch_time.append(time.time() - end)
losses.append(loss.item())
end = time.time()
# print progress
if (batch_idx + 1) % args.print_freq == 0 or (batch_idx +
1) == len(train_loader):
print('Epoch: {} [{}/{} ({:.3f}%)] Time {batch_time.last_avg:.3f}'
'\tLoss: {loss.last_avg:.4f}'.format(
epoch,
batch_idx + 1,
len(train_loader), (batch_idx + 1) / len(train_loader),
batch_time=batch_time,
loss=losses))
def train(args,
model,
device,
train_loader,
optimizer,
epoch,
descrip_str='Training'):
model.train()
pbar = tqdm(train_loader)
pbar.set_description(descrip_str)
CleanAccMeter = AvgMeter()
TradesAccMeter = AvgMeter()
for batch_idx, (data, target) in enumerate(pbar):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
# calculate robust loss
loss, cleanloss, klloss, cleanacc, tradesacc = trades_loss(
model=model,
x_natural=data,
y=target,
optimizer=optimizer,
device=device,
step_size=args.step_size,
epsilon=args.epsilon,
perturb_steps=args.num_steps,
beta=args.beta,
)
loss.backward()
optimizer.step()
CleanAccMeter.update(cleanacc)
TradesAccMeter.update(tradesacc)
pbar_dic = OrderedDict()
pbar_dic['cleanloss'] = '{:.3f}'.format(cleanloss)
pbar_dic['klloss'] = '{:.3f}'.format(klloss)
pbar_dic['CleanAcc'] = '{:.2f}'.format(CleanAccMeter.mean)
pbar_dic['TradesAcc'] = '{:.2f}'.format(TradesAccMeter.mean)
pbar.set_postfix(pbar_dic)
def train(args, model, device, train_loader, optimizer, epoch, log):
model.train()
dataTimeAve = AverageMeter()
totalTimeAve = AverageMeter()
end = time.time()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
dataTime = time.time() - end
dataTimeAve.update(dataTime)
optimizer.zero_grad()
# calculate robust loss
loss = trades_loss(model=model,
x_natural=data,
y=target,
optimizer=optimizer,
step_size=args.step_size,
epsilon=args.epsilon,
perturb_steps=args.num_steps_train,
beta=args.beta,
trainmode=args.trainmode,
fixbn=args.fixbn,
fixmode=args.fixmode)
loss.backward()
optimizer.step()
totalTime = time.time() - end
totalTimeAve.update(totalTime)
end = time.time()
# print progress
if batch_idx % args.log_interval == 0:
log.info(
'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tData time: {:.3f}\tTotal time: {:.3f}'
.format(epoch, batch_idx * len(data),
len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item(),
dataTimeAve.avg, totalTimeAve.avg))
def train(args, model, device, train_loader, optimizer, epoch):
model.train() # model.train() for training while model.learn() for testing
for batch_idx, (data, target) in enumerate(train_loader): # for i, (X,y) in enumerate(training data)
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
# calculate robust loss
loss = trades_loss(model=model,
x_natural=data,
y=target,
optimizer=optimizer,
step_size=args.step_size,
epsilon=args.epsilon,
perturb_steps=args.num_steps,
beta=args.beta, # balance for TRADES
distance='l_inf')
loss.backward() # calculate gradients
optimizer.step() # again our old friend: zero_grad -> backward -> step (update)
if batch_idx % args.log_interval == 0: # check how many times we should log in the concole our epoch num, ratio of trained data, 100% ratio, and loss
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
def train(args, model, device, train_loader, optimizer, epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
origin_output = model(data)
optimizer.zero_grad()
grad = get_grad(model, origin_output, target, optimizer,
nn.CrossEntropyLoss())
optimizer.zero_grad()
# calculate robust loss
loss = trades_loss(model=model,
x_natural=data,
y=target,
optimizer=optimizer,
step_size=args.step_size,
epsilon=args.epsilon,
perturb_steps=args.num_steps,
beta=args.beta)
loss.backward()
adv_grad = get_clear_grad(model)
# gravel
layer_mask = get_grad_diff_layer_mask(grad, adv_grad, ratio=0.1)
ret_grad = generate_grad(grad, adv_grad, layer_mask=layer_mask)
set_grad(model, ret_grad)
optimizer.step()
# print progress
if batch_idx % args.log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
def do_train_trades(model_name,
model,
train_loader,
val_loader,
adv_val_loader,
device,
i_ep,
trainStr,
lr=0.0001,
n_ep=40,
save_path='/tmp'):
optimizer = AdamW(model.parameters(), weight_decay=1e-4, lr=lr)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
factor=0.5,
verbose=True,
patience=4)
# do training
i_ep = i_ep + 1
#绕动规则
#设置对抗样本的绕动，训练时逐步加大，不要一开始就设置很大，例如可以：10训练收敛后，在设置16,20
rule_eps = [4, 6, 8, 10, 12, 14, 16, 18, 20, 22]
perturb_steps = 5
std = 0.1
# epsilon_low = 3.51
# epsilon_high = 7.02
#epsilon =0.086
print('lr is ', lr)
while i_ep < n_ep:
model.train()
train_losses = []
widgets = [
'train :',
Percentage(), ' ',
Bar('#'), ' ',
Timer(), ' ',
ETA(), ' ',
FileTransferSpeed()
]
pbar = ProgressBar(widgets=widgets)
if i_ep < len(rule_eps):
epsilon = rule_eps[i_ep] / 255
i_ep += 1
else:
epsilon = 22 / 255
print('current std is ', std)
print('current epslion is ', epsilon)
for batch_data in pbar(train_loader):
image = batch_data['image'].to(device)
label = batch_data['label_idx'].to(device)
optimizer.zero_grad()
#logits = model(image)
#loss = F.cross_entropy(logits, label)
#最核心部分,先对图片加入高斯噪声后，再用I-fgsm的方式找其对应的对抗样本,使用trade off的loss函数
#训练具有高斯噪声的模型，不能以来设很大的std,一开始设置0.1，收敛后再逐步加大继续训练
buffer = torch.Tensor(image.size()).normal_(0, std).cuda()
image = image + buffer.detach()
image = torch.clamp(image, 0.0, 1.0)
loss = trades_loss(model,
image,
label,
optimizer,
epsilon=epsilon,
perturb_steps=perturb_steps)
loss.backward()
optimizer.step()
train_losses += [loss.detach().cpu().numpy().reshape(-1)]
train_losses = np.concatenate(train_losses).reshape(-1).mean()
# 观察模型在验证集上表现
model.eval()
val_losses = []
preds = []
true_labels = []
widgets = [
'val:',
Percentage(), ' ',
Bar('#'), ' ',
Timer(), ' ',
ETA(), ' ',
FileTransferSpeed()
]
pbar = ProgressBar(widgets=widgets)
#测试时同样对图片加入高斯噪声，再前向传播，因为具有随机性，测试时也需要多次测试取logit的均值，30次左右就可稳定输出
#test_time1 是val干净图片测试次数.数量较多这里设置的10
#test_time 是攻击过的图片测试次数
test_time1 = 10
test_time = 30
for batch_data in pbar(val_loader):
image = batch_data['image'].to(device)
label = batch_data['label_idx'].to(device)
all_logits = torch.zeros((image.size()[0], 110)).to(device)
for _time in range(test_time1):
buffer = torch.Tensor(image.size()).normal_(0, std).cuda()
image_gs = image.detach() + buffer
image_gs = torch.clamp(image_gs, 0.0, 1.0)
with torch.no_grad():
logits = model(image_gs)
all_logits.add_(logits)
all_logits.div_(test_time1)
loss = F.cross_entropy(all_logits,
label).detach().cpu().numpy().reshape(-1)
val_losses += [loss]
true_labels += [label.detach().cpu().numpy()]
predictedList = all_logits.max(1)[1].detach().cpu().numpy()
preds += [(predictedList)]
preds = np.concatenate(preds, 0).reshape(-1)
true_labels = np.concatenate(true_labels, 0).reshape(-1)
acc = np.mean(true_labels == preds)
val_losses = np.concatenate(val_losses).reshape(-1).mean()
# 验证攻击acc
model.eval()
val_losses_adv = []
preds_adv = []
true_labels_adv = []
widgets2 = [
'val:',
Percentage(), ' ',
Bar('#'), ' ',
Timer(), ' ',
ETA(), ' ',
FileTransferSpeed()
]
pbar2 = ProgressBar(widgets=widgets2)
for batch_data in pbar2(adv_val_loader):
image_adv = batch_data['image'].to(device)
label_adv = batch_data['label_idx'].to(device)
all_adv_logits = torch.zeros((image_adv.size()[0], 110)).to(device)
for _time in range(test_time):
buffer = torch.Tensor(image_adv.size()).normal_(0, std).cuda()
x_adv_gs = image_adv.detach() + buffer
x_adv_gs = torch.clamp(x_adv_gs, 0.0, 1.0)
with torch.no_grad():
logits_adv = model(x_adv_gs)
all_adv_logits.add_(logits_adv)
all_adv_logits.div_(test_time)
loss_adv = F.cross_entropy(
all_adv_logits, label_adv).detach().cpu().numpy().reshape(-1)
val_losses_adv += [loss_adv]
true_labels_adv += [label_adv.detach().cpu().numpy()]
predictedList_adv = all_adv_logits.max(1)[1].detach().cpu().numpy()
preds_adv += [(predictedList_adv)]
preds_adv = np.concatenate(preds_adv, 0).reshape(-1)
true_labels_adv = np.concatenate(true_labels_adv, 0).reshape(-1)
acc_adv = np.mean(true_labels_adv == preds_adv)
val_losses_adv = np.concatenate(val_losses_adv).reshape(-1).mean()
scheduler.step(val_losses)
print(
f'Epoch : {i_ep} val_acc : {acc:.5%} ||| val_adc_acc :{acc_adv:.5%} ||| train_loss : {train_losses:.5f} val_loss : {val_losses:.5f} ||| val_loss_adv:{val_losses_adv:.5f}|||'
)
#每个周期都保存
torch.save({
"model_state_dict": model.state_dict(),
"i_ep": i_ep,
},
os.path.join(
save_path,
f'ep_{i_ep}_{model_name}_val_acc_{acc:.4f}.pth'))
model.to(device)
i_ep += 1
def train(args, model, device, train_loader, optimizer, epoch):
model.train()
train_metrics = []
unlabeled_robust_weight = adjust_unlabeled_robust_weight(epoch)
epsilon = adjust_epsilon(epoch)
#logger.info('Using unlabeled robust weight of %g' % unlabeled_robust_weight)
for batch_idx, (data, target, unsup) in enumerate(train_loader):
data, target, unsup = data.to(device), target.to(device), unsup.to(
device)
optimizer.zero_grad()
# calculate robust loss
if args.loss == 'trades':
(loss, natural_loss, robust_loss, robust_loss_labeled,
entropy_loss_unlabeled) = trades_loss(
model=model,
epoch=epoch,
unsup=unsup,
x_natural=data,
y=target,
optimizer=optimizer,
step_size=args.step_size,
epsilon=epsilon,
perturb_steps=args.num_steps,
beta=args.beta,
distance=args.distance,
batch_mode=args.batch_mode,
extra_class_weight=args.unlabeled_class_weight,
entropy_weight=args.entropy_weight,
unlabeled_natural_weight=args.unlabeled_natural_weight,
unlabeled_robust_weight=unlabeled_robust_weight,
init=args.attack_init)
elif args.loss == 'madry':
assert (not args.semisup
), 'Cannot use unsupervised data with Madry loss'
assert (args.beta <= 1), 'Madry loss requires beta < 1'
(loss, natural_loss, robust_loss, loss_natural_labeled,
loss_natural_unlabeled) = madry_loss(
model=model,
epoch=epoch,
x_natural=data,
y=target,
unsup=unsup,
optimizer=optimizer,
step_size=args.step_size,
epsilon=args.epsilon,
perturb_steps=args.num_steps,
beta=args.beta,
unlabeled_natural_weight=args.unlabeled_natural_weight,
unlabeled_robust_weight=args.unlabeled_robust_weight,
distance=args.distance,
batch_mode=args.batch_mode,
extra_class_weight=args.unlabeled_class_weight,
max_rot=args.max_rot,
max_trans=args.max_trans)
robust_loss_labeled = loss_natural_labeled
entropy_loss_unlabeled = loss_natural_unlabeled
loss.backward()
optimizer.step()
curr_train_metrics = dict(
epoch=epoch,
loss=loss.item(),
natural_loss=natural_loss.item(),
robust_loss=robust_loss.item(),
robust_loss_labeled=robust_loss_labeled.item(),
entropy_loss_unlabeled=entropy_loss_unlabeled.item())
train_metrics.append(curr_train_metrics)
# print progress
if batch_idx % args.log_interval == 0:
logging.info(
'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), epoch_size,
100. * batch_idx / len(train_loader), loss.item()))
return train_metrics
