Function eval_on_batch

script/feature/direct_feature_matching.py:178–213

Perform 1 step of eval

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

    return batch_rays, target_s

def eval_on_batch(args, data, model, feat_model, pose, img_idx, hwf, half_res, device, world_setup_dict, **render_kwargs_test):
    ''' Perform 1 step of eval'''
    with torch.no_grad():
        H, W, focal = hwf
        _, 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, world_setup_dict, device=device_cpu)

        half_res=False # no need to use half_res for inference
        batch_rays, target = prepare_batch_render(args, pose_nerf, args.batch_size, data, 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)
        pose_nerf = pose_nerf.to(device)

        # every new tensor from onward is in GPU
        torch.set_default_tensor_type('torch.cuda.FloatTensor')
        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, feat_model, hwf, half_res, device, world_setup_dict, **render_kwargs_test):
    ''' Perform 1 epoch of training with batch '''