def obtain_predictions(model, device, loader, known_depth=False, print_progress=False): model.eval() iterable = loader if print_progress: iterable = tqdm(loader, leave=True, ascii=True) for batch in iterable: in_var = batch['input'].to(device, torch.float32) target_var = batch['target'].to(device, torch.float32) # Calculate predictions and loss start_time = perf_counter() out_var = model(in_var) inference_time = perf_counter() - start_time loss = average_loss( model.forward_3d_losses(out_var, target_var.narrow(-1, 0, 3))) norm_preds = ensure_homogeneous(out_var.to(CPU, torch.float64), d=3) actuals = [] expected = None for i, norm_pred in enumerate(norm_preds): expected_i, actual_i =\ prepare_for_3d_evaluation(batch['original_skel'][i], norm_pred, loader.dataset, batch['camera_intrinsic'][i], batch['transform_opts'][i], known_depth=known_depth) if expected is not None: assert (expected_i - expected).abs().gt(1e-6).sum() == 0,\ "Expected all examples in batch to have the same target" expected = expected_i actuals.append(actual_i) actual = torch.stack(actuals, 0).mean(0) try: frame_ref = batch['frame_ref'][0] except KeyError: frame_ref = None prediction = dict( expected=expected, actual=actual, frame_ref=frame_ref, inference_time=inference_time, loss=loss.sum().item(), ) yield prediction
def forward_loss(model, out_var, target_var, mask_var, valid_depth): target_var = target_var.narrow(-1, 0, 3) if not 0 in valid_depth: losses = model.forward_3d_losses(out_var, target_var) elif not 1 in valid_depth: losses = model.forward_2d_losses(out_var, target_var) else: losses_3d = model.forward_3d_losses(out_var, target_var) losses_2d = model.forward_2d_losses(out_var, target_var) losses = torch.stack([(losses_3d[i] if use_3d == 1 else losses_2d[i]) for i, use_3d in enumerate(valid_depth)]) return average_loss(losses, mask_var)