def all_reduce_tensor(self, tensor, norm=True):
if self.distributed:
return all_reduce_tensor(tensor,
world_size=self.world_size,
norm=norm)
else:
return torch.mean(tensor)
cur_masks = cur_masks.cuda(non_blocking=True)
preds = model(ref_imgs, cur_imgs)
preds = F.interpolate(preds,
(cur_masks.size()[2], cur_masks.size()[3]),
mode='bilinear',
align_corners=False)
loss = criterion(preds, cur_masks)
#loss = lovasz_hinge(preds.squeeze(), cur_masks.squeeze())
loss = bootstrapped_ce_loss(loss)
loss += 1. * dice_loss(preds[:, 0].flatten(1),
cur_masks.flatten(1))
# reduce the whole loss over multi-gpu
if engine.distributed:
reduce_loss = all_reduce_tensor(loss,
world_size=engine.world_size)
optimizer.zero_grad()
loss.backward()
optimizer.step()
current_idx = epoch * config.niters_per_epoch + idx
lr = lr_policy.get_lr(current_idx)
optimizer.param_groups[0]['lr'] = lr
optimizer.param_groups[1]['lr'] = lr
for i in range(2, len(optimizer.param_groups)):
optimizer.param_groups[i]['lr'] = lr * 10
writer.add_scalar("train/loss",
scalar_value=reduce_loss.item(),
engine.update_iteration(epoch, idx)
minibatch = dataloader.next()
imgs = minibatch['data']
gts = minibatch['label']
edge_gts = minibatch['aux_label']
imgs = imgs.cuda(non_blocking=True)
gts = gts.cuda(non_blocking=True)
edge_gts = edge_gts.cuda(non_blocking=True)
loss, edge_loss, onelevel_segloss, twolevel_segloss = model(
imgs, gts, edge_gts)
# reduce the whole loss over multi-gpu
if engine.distributed:
reduce_loss = all_reduce_tensor(loss,
world_size=engine.world_size)
reduce_edge_loss = all_reduce_tensor(
edge_loss, world_size=engine.world_size)
onelevel_segloss = all_reduce_tensor(
onelevel_segloss, world_size=engine.world_size)
twolevel_segloss = all_reduce_tensor(
twolevel_segloss, world_size=engine.world_size)
optimizer.zero_grad()
loss.backward()
optimizer.step()
current_idx = epoch * config.niters_per_epoch + idx
lr = lr_policy.get_lr(current_idx)
optimizer.param_groups[0]['lr'] = lr
