def train_one_epoch(model, dataloader_source, dataloader_target, epoch, args):
model.train()
len_dataloader = min(len(dataloader_source), len(dataloader_target))
data_source_iter = iter(dataloader_source)
data_target_iter = iter(dataloader_target)
i = 0
while i + 1 < len_dataloader:
# training model using source data
s_img, s_label = data_source_iter.next()
s_img = s_img.expand(s_img.data.shape[0], 3, 28, 28)
s_img = s_img.cuda()
s_label = s_label.cuda()
# training model using target data
t_img, t_label = data_target_iter.next()
t_img = t_img.expand(t_img.data.shape[0], 3, 28, 28)
t_img = t_img.cuda()
t_label = t_label.cuda()
# s_class_output= model(input_data=s_img)
# t_class_output = model(input_data=t_img)
#
# loss_s = CEL(s_class_output, s_label)
# robust_loss_t = CEL(t_class_output, t_label) + beta * KLL(F.log_softmax(t_class_output, dim=1), F.softmax(s_class_output, dim=1))
#
# loss = loss_s + weight * robust_loss_t
cat_img = torch.cat((s_img, t_img), 0)
cat_target = torch.cat((s_label, t_label), 0)
(loss, natural_loss, robust_loss,
entropy_loss_unlabeled) = trades_loss(
model=model,
x_natural=cat_img,
y=cat_target,
optimizer=optimizer,
step_size=args.pgd_step_size,
epsilon=args.epsilon,
perturb_steps=args.pgd_num_steps,
beta=args.beta,
distance=args.distance,
adversarial=args.distance == 'l_inf',
entropy_weight=args.entropy_weight)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# log to tensorboard
global_step = epoch * (len_dataloader - 1) + i
logger.add_scalar("loss", loss, global_step)
i += 1
if i % 100 == 0:
print('epoch: %d, [iter: %d / all %d], loss: %f' \
% (epoch, i, len_dataloader, loss.cpu().data.numpy()))
def train(args, model, device, train_loader, optimizer, epoch):
model.train()
train_metrics = []
epsilon = args.epsilon
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
# calculate robust loss
if args.loss == 'trades':
# The TRADES KL-robustness regularization term proposed by
# Zhang et al., with some additional features
(loss, natural_loss, robust_loss,
entropy_loss_unlabeled) = trades_loss(
model=model,
x_natural=data,
y=target,
optimizer=optimizer,
step_size=args.pgd_step_size,
epsilon=epsilon,
perturb_steps=args.pgd_num_steps,
beta=args.beta,
distance=args.distance,
adversarial=args.distance == 'l_inf',
entropy_weight=args.entropy_weight)
elif args.loss == 'noise':
# Augmenting the input with random noise as in Cohen et al.
assert (args.distance == 'l_2')
loss = noise_loss(model=model,
x_natural=data,
y=target,
clamp_x=True,
epsilon=epsilon)
entropy_loss_unlabeled = torch.Tensor([0.])
natural_loss = robust_loss = loss
loss.backward()
optimizer.step()
train_metrics.append(
dict(epoch=epoch,
loss=loss.item(),
natural_loss=natural_loss.item(),
robust_loss=robust_loss.item(),
entropy_loss_unlabeled=entropy_loss_unlabeled.item()))
# print progress
if batch_idx % args.log_interval == 0:
logging.info(
'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, min(batch_idx * args.batch_size, epoch_size),
epoch_size, 100. * batch_idx / len(train_loader),
loss.item()))
return train_metrics
def train(self, epoch):
lr = adjust_learning_rate(self.lr, self.optimizer, self.config.lr_d,
epoch)
self.model.train()
losses = AveMeter()
f_nfe_meter = AveMeter()
b_nfe_meter = AveMeter()
correct = 0
for i, (inputs, targets) in enumerate(self.loaders['train']):
inputs = inputs + self.config.noise * torch.randn_like(inputs)
if self.config.use_gpu:
inputs, targets = inputs.cuda(), targets.cuda()
self.optimizer.zero_grad()
if self.is_ode:
self.model.module.nfe = 0
outputs = self.model(inputs)
if self.is_ode:
nfe_forward = self.model.module.nfe.item()
if self.config.loss == 'ce':
loss = self.criterion(outputs, targets)
elif self.config.loss == 'tr':
loss = self.criterion(outputs, targets) + trades_loss(
self.model,
inputs,
self.optimizer,
distance=self.config.attack_type)
elif self.config.loss == 'ma':
loss = madry_loss(self.model, inputs, targets, self.optimizer)
if self.is_ode:
self.model.module.nfe = 0
loss.backward()
self.optimizer.step()
if self.is_ode:
nfe_backward = self.model.module.nfe.item()
self.model.module.nfe = 0
if self.is_ode:
f_nfe_meter.update(nfe_forward)
b_nfe_meter.update(nfe_backward)
losses.update(loss.item(), inputs.size()[0])
correct += outputs.max(1)[1].eq(targets).sum().item()
acc = 100 * correct / len(self.datasets['train'])
logger.info(
f"Train: [{i + 1}/{len(self.loaders['train'])}] | "
f"Time: {self.timer.timeSince()} | "
f"loss: {losses.avg:.4f} | "
f"acc: {acc:.2f}% | NFE-F: {f_nfe_meter.avg:.2f} | NFE-B: {b_nfe_meter.avg:.2f}"
)
self.writer.add_scalar('train/lr', lr, epoch)
self.writer.add_scalar('train/loss', losses.avg, epoch)
self.writer.add_scalar('train/acc', acc, epoch)
self.writer.add_scalar('train/nfe-f', f_nfe_meter.avg, epoch)
self.writer.add_scalar('train/nfe-b', b_nfe_meter.avg, epoch)
for name, param in self.model.named_parameters():
if param.grad is not None:
self.writer.add_histogram(name,
param.clone().cpu().data.numpy(),
epoch)
self.writer.add_histogram(
name + '/grad',
param.grad.clone().cpu().data.numpy(), epoch)