def save_model_predictions(p, val_loader, model):
""" Save model predictions for all tasks """
print('Save model predictions to {}'.format(p['save_dir']))
model.eval()
tasks = p.TASKS.NAMES
save_dirs = {task: os.path.join(p['save_dir'], task) for task in tasks}
for save_dir in save_dirs.values():
mkdir_if_missing(save_dir)
for ii, sample in enumerate(val_loader):
inputs, meta = sample['image'].cuda(non_blocking=True), sample['meta']
img_size = (inputs.size(2), inputs.size(3))
output = model(inputs)
for task in p.TASKS.NAMES:
output_task = get_output(output[task], task).cpu().data.numpy()
for jj in range(int(inputs.size()[0])):
if len(sample[task][jj].unique()
) == 1 and sample[task][jj].unique() == 255:
continue
fname = meta['image'][jj]
result = cv2.resize(output_task[jj],
dsize=(meta['im_size'][1][jj],
meta['im_size'][0][jj]),
interpolation=p.TASKS.INFER_FLAGVALS[task])
if task == 'depth':
sio.savemat(os.path.join(save_dirs[task], fname + '.mat'),
{'depth': result})
else:
imageio.imwrite(
os.path.join(save_dirs[task], fname + '.png'),
result.astype(np.uint8))
def eval_model(p, val_loader, model):
""" Evaluate model in an online fashion without storing the predictions to disk """
tasks = p.TASKS.NAMES
performance_meter = PerformanceMeter(p)
model.eval()
for i, batch in enumerate(val_loader):
# Forward pass
images = batch['image'].cuda(non_blocking=True)
targets = {task: batch[task].cuda(non_blocking=True) for task in tasks}
output = model(images)
# Measure performance
performance_meter.update({t: get_output(output[t], t) for t in tasks}, targets)
eval_results = performance_meter.get_score(verbose = True)
return eval_results
def train_vanilla(p, train_loader, model, criterion, optimizer, epoch):
""" Vanilla training with fixed loss weights """
losses = get_loss_meters(p)
performance_meter = PerformanceMeter(p)
progress = ProgressMeter(len(train_loader), [v for v in losses.values()],
prefix="Epoch: [{}]".format(epoch))
model.train()
for i, batch in enumerate(train_loader):
# Forward pass
images = batch['image'].cuda(non_blocking=True)
targets = {
task: batch[task].cuda(non_blocking=True)
for task in p.ALL_TASKS.NAMES
}
output = model(images)
# Measure loss and performance
loss_dict = criterion(output, targets)
for k, v in loss_dict.items():
losses[k].update(v.item())
performance_meter.update(
{t: get_output(output[t], t)
for t in p.TASKS.NAMES}, {t: targets[t]
for t in p.TASKS.NAMES})
# Backward
optimizer.zero_grad()
loss_dict['total'].backward()
optimizer.step()
if i % 25 == 0:
progress.display(i)
eval_results = performance_meter.get_score(verbose=True)
return eval_results