Function train_on_batch

script/dm/direct_pose_model.py:228–276

Perform 1 step of training

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

    return total_loss_mean, total_psnr_mean

def train_on_batch(args, data, model, pose, img_idx, hwf, optimizer, half_res, device, **render_kwargs_test):
    ''' Perform 1 step of training'''
    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 Design Here ###
    # Compute RGB MSE Loss
    photo_loss = rgb_loss(rgb, target, extras)

    # Compute Combine Loss if needed
    if args.combine_loss:
        pose_loss = PoseLoss(args, pose_, pose, device)
        loss = args.combine_loss_w[0] * pose_loss + args.combine_loss_w[1] * photo_loss
        # print("img_idx: {}, pose_loss: {}, rgb_loss {}, total_loss {}".format(img_idx, args.combine_loss_w[0] * pose_loss, args.combine_loss_w[1] * photo_loss, loss))

    ### Loss Design End
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    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 train_on_epoch(args, data_loaders, model, hwf, optimizer, half_res, device, **render_kwargs_test):