def train(train_loader, plfd_backbone, criterion, optimizer, epoch):
losses = AverageMeter()
plfd_backbone.train(True)
for img, landmark_gt, attribute_gt, euler_angle_gt in train_loader:
img.requires_grad = False
img = img.cuda()
optimizer.zero_grad()
attribute_gt.requires_grad = False
attribute_gt = attribute_gt.cuda()
landmark_gt.requires_grad = False
landmark_gt = landmark_gt.cuda()
euler_angle_gt.requires_grad = False
euler_angle_gt = euler_angle_gt.cuda()
plfd_backbone = plfd_backbone.cuda()
pose, landmarks = plfd_backbone(img)
# attribute_gt, landmark_gt, euler_angle_gt, angle, landmarks, train_batchsize)
lds_loss, pose_loss = criterion(attribute_gt, landmark_gt,
euler_angle_gt, pose, landmarks,
args.train_batchsize)
# pose_loss = pose_criterion(pose, euler_angle_gt)
# lds_loss = pose_criterion(landmarks, landmark_gt)
total_loss = pose_loss + lds_loss
total_loss.backward()
optimizer.step()
losses.update(total_loss.item())
return pose_loss, lds_loss
def train(train_loader, plfd_backbone, auxiliarynet, criterion, optimizer,
epoch):
losses = AverageMeter()
for img, landmark_gt, attribute_gt, euler_angle_gt in train_loader:
img.requires_grad = False
img = img.cuda(non_blocking=True)
attribute_gt.requires_grad = False
attribute_gt = attribute_gt.cuda(non_blocking=True)
landmark_gt.requires_grad = False
landmark_gt = landmark_gt.cuda(non_blocking=True)
euler_angle_gt.requires_grad = False
euler_angle_gt = euler_angle_gt.cuda(non_blocking=True)
plfd_backbone = plfd_backbone.cuda()
auxiliarynet = auxiliarynet.cuda()
features, landmarks = plfd_backbone(img)
angle = auxiliarynet(features)
weighted_loss, loss = criterion(attribute_gt, landmark_gt,
euler_angle_gt, angle, landmarks,
args.train_batchsize)
optimizer.zero_grad()
weighted_loss.backward()
optimizer.step()
losses.update(loss.item())
return weighted_loss, loss
def train(train_loader, plfd_backbone, auxiliarynet, criterion, optimizer,
epoch):
losses = AverageMeter()
iter_num = 0
for img, landmark_gt, attribute_gt, euler_angle_gt in train_loader:
start = time.time()
img = img.to(device)
attribute_gt = attribute_gt.to(device)
landmark_gt = landmark_gt.to(device)
euler_angle_gt = euler_angle_gt.to(device)
plfd_backbone = plfd_backbone.to(device)
auxiliarynet = auxiliarynet.to(device)
features, landmarks = plfd_backbone(img)
angle = auxiliarynet(features)
weighted_loss, loss = criterion(attribute_gt, landmark_gt,
euler_angle_gt, angle, landmarks,
args.train_batchsize)
optimizer.zero_grad()
weighted_loss.backward()
optimizer.step()
time_cost_per_batch = time.time() - start
losses.update(loss.item())
iter_num += 1
if iter_num % 10 == 0:
msg = "Epoch: {}, Iter: {}, Loss: {:.6f}, Weight_Loss: {:.6f}, Speed: {} imgs/sec".format(
epoch, iter_num, loss.item(), weighted_loss.item(),
256 // time_cost_per_batch)
logging.info(msg)
print("===> Train:")
print('Epoch: {}, Average loss: {:.6f} '.format(epoch, losses.avg))
return weighted_loss, loss
def train(train_loader, plfd_backbone, auxiliarynet, criterion, optimizer,
epoch):
losses = AverageMeter()
for img, landmark_gt, attribute_gt, euler_angle_gt in train_loader:
img = img.to(device)
attribute_gt = attribute_gt.to(device)
landmark_gt = landmark_gt.to(device)
euler_angle_gt = euler_angle_gt.to(device)
plfd_backbone = plfd_backbone.to(device)
auxiliarynet = auxiliarynet.to(device)
plfd_backbone.apply(lambda m: type(m) == AdvNoise and m.set_clean())
optimizer.zero_grad()
features, landmarks = plfd_backbone(img)
angle = auxiliarynet(features)
weighted_loss, loss = criterion(attribute_gt, landmark_gt,
euler_angle_gt, angle, landmarks,
args.train_batchsize)
weighted_loss.backward()
plfd_backbone.apply(lambda m: type(m) == AdvNoise and m.set_stay())
optimizer.zero_grad()
features, landmarks = plfd_backbone(img)
angle = auxiliarynet(features)
weighted_loss, loss = criterion(attribute_gt, landmark_gt,
euler_angle_gt, angle, landmarks,
args.train_batchsize)
weighted_loss.backward()
optimizer.step()
losses.update(loss.item())
return weighted_loss, loss
def train(train_loader, plfd_backbone, auxiliarynet, criterion, optimizer,
epoch):
losses = AverageMeter()
num_batch = len(train_loader)
i = 0
for img, landmark_gt, attribute_gt, euler_angle_gt in train_loader:
i += 1
img = img.to(device)
attribute_gt = attribute_gt.to(device)
landmark_gt = landmark_gt.to(device)
euler_angle_gt = euler_angle_gt.to(device)
plfd_backbone = plfd_backbone.to(device)
auxiliarynet = auxiliarynet.to(device)
features, landmarks = plfd_backbone(img)
angle = auxiliarynet(features)
using_wingloss = wandb.config.using_wingloss
weighted_loss, loss = criterion(attribute_gt, landmark_gt, euler_angle_gt,
angle, landmarks, args.train_batchsize, using_wingloss=using_wingloss)
optimizer.zero_grad()
weighted_loss.backward()
optimizer.step()
losses.update(loss.item())
wandb.log({"metric/loss":loss.item()})
wandb.log({"metric/weighted_loss": weighted_loss.detach().numpy()})
logger.info(f"Epoch:{epoch}. Batch {i} / {num_batch} batches. Loss: {loss.item()}. Weighted_loss:{ weighted_loss.detach().numpy()}")
return weighted_loss, loss
def train(train_loader, pfld_backbone, auxiliarynet, criterion, optimizer,
epoch):
losses = AverageMeter()
weighted_loss, loss = None, None
for img, landmark_gt, attribute_gt, euler_angle_gt in train_loader:
img = img.to(device)
attribute_gt = attribute_gt.to(device)
landmark_gt = landmark_gt.to(device)
euler_angle_gt = euler_angle_gt.to(device)
pfld_backbone = pfld_backbone.to(device)
auxiliarynet = auxiliarynet.to(device)
features, landmarks = pfld_backbone(img)
angle = auxiliarynet(features)
weighted_loss, loss = criterion(attribute_gt, landmark_gt,
euler_angle_gt, angle, landmarks,
args.train_batchsize)
optimizer.zero_grad()
weighted_loss.backward()
optimizer.step()
losses.update(loss.item())
return weighted_loss, loss
def train(train_loader,
plfd_backbone,
auxiliarynet,
criterion,
optimizer,
epoch,
log_interval=10):
losses = AverageMeter()
plfd_backbone.train()
auxiliarynet.train()
# print('is_training:', plfd_backbone.training)
logging.info("total iteration is {}".format(len(train_loader)))
weighted_loss, loss = 0, 0
for iteration, (img, landmark_gt,
euler_angle_gt) in enumerate(train_loader):
img = img.to(device)
landmark_gt = landmark_gt.to(device)
euler_angle_gt = euler_angle_gt.to(device)
plfd_backbone = plfd_backbone.to(device)
auxiliarynet = auxiliarynet.to(device)
features, landmarks = plfd_backbone(img)
angle = auxiliarynet(features)
weighted_loss, loss = criterion(landmark_gt, euler_angle_gt, angle,
landmarks, args.train_batchsize)
# weighted_loss, loss = criterion(landmarks, landmark_gt)
optimizer.zero_grad()
weighted_loss.backward()
optimizer.step()
losses.update(loss.item())
if iteration % log_interval == 0:
logging.info("epoch: {}, iteration: {}, train loss: {:.4f}".format(
epoch, iteration, weighted_loss.item()))
return weighted_loss, loss