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