Function train_on_feature_batch

script/feature/direct_feature_matching.py:235–320

Perform 1 step of training using scheme1

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

    return total_loss_mean, total_psnr_mean

def train_on_feature_batch(args, data, model, feat_model, pose, img_idx, hwf, optimizer, device, world_setup_dict, **render_kwargs_test):
    ''' Perform 1 step of training using scheme1 '''
    batch_size_iter = data.shape[0]

    H, W, focal = hwf
    data = data.to(device) # [1, 3, 240, 427]
    
    # pose regression module
    _, pose_ = inference_pose_regression(args, data, device, model, retFeature=False) # here returns predicted pose [1, 3, 4] # real img features and predicted pose # features: (1, [3, 1, 128, 240, 427]), predict_pose: [1, 3, 4]
    pose_nerf = pose_.clone()

    # rescale the predicted pose to nerf scales
    pose_nerf = fix_coord_supp(args, pose_nerf, world_setup_dict, device=device)

    pose = pose.to(device)
    img_idx = img_idx.to(device)
    # every new tensor from onward is in GPU, here memory cost is a bottleneck
    torch.set_default_tensor_type('torch.cuda.FloatTensor')

    # here is single frame
    target = data.permute(0,2,3,1) # [B,H,W,C]
    rays_o_list=[]
    rays_d_list=[]
    img_idx_list=[]
    N_rand = args.N_rand
    for i in range(pose_nerf.shape[0]):
        rays_o, rays_d = get_rays(H, W, focal, pose_nerf[i])  # (H, W, 3), (H, W, 3)
        rays_o_list.append(rays_o)
        rays_d_list.append(rays_d)
        img_idx_list.append(img_idx[i].repeat(N_rand,1))
    rays_o_batch = torch.stack(rays_o_list)
    rays_d_batch = torch.stack(rays_d_list)
    img_idx_batch = torch.cat(img_idx_list)

    # randomly select coords
    coords = torch.stack(torch.meshgrid(torch.linspace(0, H-1, H), torch.linspace(0, W-1, W), indexing='ij'), -1)  # (H, W, 2)
    coords = torch.reshape(coords, [-1,2])  # (H * W, 2)
    select_inds = np.random.choice(coords.shape[0], size=[N_rand], replace=False)  # (N_rand,)
    select_coords = coords[select_inds].long()  # (N_rand, 2)

    # fetch from coords
    rays_o = rays_o_batch[:, select_coords[:, 0], select_coords[:, 1]]
    rays_d = rays_d_batch[:, select_coords[:, 0], select_coords[:, 1]]
    rays_o = rays_o.reshape(rays_o.shape[0]*rays_o.shape[1], 3) # (B*N_rand, 3)
    rays_d = rays_d.reshape(rays_d.shape[0]*rays_d.shape[1], 3) # (B*N_rand, 3)
    batch_rays = torch.stack([rays_o, rays_d], 0)
    target_s = target[:,select_coords[:, 0], select_coords[:, 1]].reshape(batch_size_iter*N_rand,3)  # (B*N_rand, 3)

    rgb_feature, disp, acc, extras = render(H, W, focal, chunk=args.chunk, rays=batch_rays, img_idx=img_idx_batch, **render_kwargs_test)
    # rgb_feature is rgb 3 + features 128
    rgb = rgb_feature[...,:3] # [B*N_rand, 3]
    feature = rgb_feature[...,3:].reshape(batch_size_iter, N_rand, args.out_channel_size-3)[None, ...].permute(0,1,3,2) # [lvl, B, C, N_rand] assuming lvl size = 1

    # inference featurenet
    target_in = target.permute(0,3,1,2)
    features, _ = feat_model(target_in, True, True, H, W) # features: (1, [3,B,C,H,W])

     # get features_target, # now choose 1st level feature only