Function eval_on_batch

script/dm/direct_pose_model.py:169–206

Perform 1 step of eval

(args, data, model, pose, img_idx, hwf, half_res, device, **render_kwargs_test)

    return pose

def eval_on_batch(args, data, model, pose, img_idx, hwf, half_res, device, **render_kwargs_test):
    ''' Perform 1 step of eval'''
    with torch.no_grad():
        H, W, focal = hwf
        target_ = deepcopy(data)
        pose_ = inference_pose_regression(args, data, device, model)
        device_cpu = torch.device('cpu')
        pose_ = pose_.to(device_cpu) # put predict pose back to cpu
        pose_nerf = pose_.clone()

        if args.NeRFH:
            # rescale the predicted pose to nerf scales
            pose_nerf = fix_coord_supp(args, pose_nerf)

        batch_rays, target = prepare_batch_render(args, pose_nerf, args.batch_size, target_, H, W, focal, half_res)
        batch_rays = batch_rays.to(device)
        target = target.to(device)
        pose = pose.to(device)
        img_idx = img_idx.to(device)

        # every new tensor from onward is in GPU
        torch.set_default_tensor_type('torch.cuda.FloatTensor')
        if half_res:
            rgb, disp, acc, extras = render(H//2, W//2, focal/2, chunk=args.chunk, rays=batch_rays, img_idx=img_idx, **render_kwargs_test)
        else:
            rgb, disp, acc, extras = render(H, W, focal, chunk=args.chunk, rays=batch_rays, img_idx=img_idx, **render_kwargs_test)

        loss = PoseLoss(args, pose_, pose, device)
        psnr = mse2psnr(img2mse(rgb, target))

        # end of every new tensor from onward is in GPU
        torch.set_default_tensor_type('torch.FloatTensor')

        iter_loss = loss.to(device_cpu).detach().numpy()
        iter_loss = np.array([iter_loss])

        iter_psnr = psnr.to(device_cpu).detach().numpy()
    return iter_loss, iter_psnr

def eval_on_epoch(args, data_loaders, model, hwf, half_res, device, **render_kwargs_test):
    ''' Perform 1 epoch of training with batch '''