minibatch = dataloader.next()
imgs = minibatch['data']
gts = minibatch['label']
imgs = imgs.cuda(non_blocking=True)
gts = gts.cuda(non_blocking=True)
loss = model(imgs, gts)
# reduce the whole loss over multi-gpu
if engine.distributed:
dist.all_reduce(loss, dist.ReduceOp.SUM)
loss = loss / engine.world_size
else:
loss = Reduce.apply(*loss) / len(loss)
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
loss.backward()
optimizer.step()
print_str = 'Epoch{}/{}'.format(epoch, config.nepochs) \
+ ' Iter{}/{}:'.format(idx + 1, config.niters_per_epoch) \
+ ' lr=%.2e' % lr \
+ ' loss=%.2f' % loss.item()
imgs = imgs.cuda(non_blocking=True)
gts = gts.cuda(non_blocking=True)
hha = hha.cuda(non_blocking=True)
aux_rate = 0.2
loss, loss_aux = model(imgs, hha, gts)
tot_loss = loss + loss_aux * config.aux_rate
# reduce the whole loss over multi-gpu
if engine.distributed:
dist.all_reduce(tot_loss, dist.ReduceOp.SUM)
tot_loss = tot_loss / engine.world_size
dist.all_reduce(loss_aux, dist.ReduceOp.SUM)
tot_loss_aux = loss_aux / engine.world_size
else:
tot_loss = Reduce.apply(*tot_loss) / len(tot_loss)
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
tot_loss.backward()
optimizer.step()
print_str = 'Epoch{}/{}'.format(epoch, config.nepochs) \
+ ' Iter{}/{}:'.format(idx + 1, config.niters_per_epoch) \
+ ' lr=%.2e' % lr \
+ ' loss=%.2f' % tot_loss.item() \
