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)