Function train_on_batch_with_random_view_synthesis

script/run_feature.py:166–230

we implement random view synthesis for generating more views to help training posenet

(args, targets, rgbs, poses, virtue_view, poses_perturb, feat_model, dset_size, FeatureLoss, optimizer, hwf, img_idxs, render_kwargs_test)

    return train_loss

def train_on_batch_with_random_view_synthesis(args, targets, rgbs, poses, virtue_view, poses_perturb, feat_model, dset_size, FeatureLoss, optimizer, hwf, img_idxs, render_kwargs_test):
    ''' we implement random view synthesis for generating more views to help training posenet '''
    feat_model.train()

    H, W, focal = hwf
    H, W = int(H), int(W)

    if args.freezeBN:
        feat_model = freeze_bn_layer_train(feat_model)

    train_loss_epoch = []

    # random generate batch_size of idx
    select_inds = np.random.choice(dset_size, size=[dset_size], replace=False)  # (N_rand,)

    batch_size=args.featurenet_batch_size # manual setting, use smaller batch size like featurenet_batch_size = 4 if OOM
    if dset_size % batch_size == 0:
        N_iters = dset_size//batch_size
    else:
        N_iters = dset_size//batch_size + 1
    
    i_batch = 0
    for i in range(0, N_iters):
        if i_batch + batch_size > dset_size:
            i_batch = 0
            break
        i_inds = select_inds[i_batch:i_batch+batch_size]
        i_batch = i_batch + batch_size

        # convert input shape to [B, 3, H, W]
        target_in = targets[i_inds].clone().permute(0,3,1,2).to(device)
        rgb_in = rgbs[i_inds].clone().permute(0,3,1,2).to(device)
        pose = poses[i_inds].clone().reshape(batch_size, 12).to(device)
        rgb_perturb = virtue_view[i_inds].clone().permute(0,3,1,2).to(device)
        pose_perturb = poses_perturb[i_inds].clone().reshape(batch_size, 12).to(device)

        # inference feature model for GT and nerf image
        pose = torch.cat([pose, pose]) # double gt pose tensor
        features, predict_pose = feat_model(torch.cat([target_in, rgb_in]), return_feature=True, upsampleH=H, upsampleW=W) # features: (1, [2, B, C, H, W])

        # get features_target and features_rgb
        if args.DFNet:
            features_target = features[0] # [3, B, C, H, W]
            features_rgb = features[1]

        loss_pose = PoseLoss(args, predict_pose, pose, device) # target

        if args.tripletloss:
            loss_f = triplet_loss_hard_negative_mining_plus(features_rgb, features_target, margin=args.triplet_margin)
        else:
            loss_f = FeatureLoss(features_rgb, features_target) # feature Maybe change to s2d-ce loss

        # inference model for RVS image
        _, virtue_pose = feat_model(rgb_perturb.to(device), False)

        # add relative pose loss here. TODO: This FeatureLoss is nn.MSE. Should be fixed later
        loss_pose_perturb = PoseLoss(args, virtue_pose, pose_perturb, device)
        loss = args.combine_loss_w[0]*loss_pose + args.combine_loss_w[1]*loss_f + args.combine_loss_w[2]*loss_pose_perturb