Perform 1 step of training
(args, data, model, feat_model, pose, img_idx, hwf, optimizer, half_res, device, world_setup_dict, **render_kwargs_test)
| 320 | return iter_loss, iter_psnr |
| 321 | |
| 322 | def train_on_batch(args, data, model, feat_model, pose, img_idx, hwf, optimizer, half_res, device, world_setup_dict, **render_kwargs_test): |
| 323 | ''' Perform 1 step of training ''' |
| 324 | |
| 325 | H, W, focal = hwf |
| 326 | data = data.to(device) # [1, 3, 240, 427] non_blocking=True |
| 327 | |
| 328 | # pose regression module |
| 329 | _, pose_ = inference_pose_regression(args, data, device, model, retFeature=False) |
| 330 | pose_nerf = pose_.clone() |
| 331 | |
| 332 | # direct matching module |
| 333 | # rescale the predicted pose to nerf scales |
| 334 | pose_nerf = fix_coord_supp(args, pose_nerf, world_setup_dict, device=device) |
| 335 | |
| 336 | pose = pose.to(device) |
| 337 | img_idx = img_idx.to(device) |
| 338 | # every new tensor from onward is in GPU, here memory cost is a bottleneck |
| 339 | torch.set_default_tensor_type('torch.cuda.FloatTensor') |
| 340 | |
| 341 | if half_res: |
| 342 | rgb, disp, acc, extras = render(H//4, W//4, focal/4, chunk=args.chunk, c2w=pose_nerf[0,:3,:4], img_idx=img_idx, **render_kwargs_test) |
| 343 | # convert rgb to B,C,H,W |
| 344 | rgb = rgb[None,...].permute(0,3,1,2) |
| 345 | # upsample rgb to hwf size |
| 346 | rgb = torch.nn.Upsample(size=(H, W), mode='bicubic')(rgb) |
| 347 | # # convert rgb back to H,W,C format |
| 348 | # rgb = rgb[0].permute(1,2,0) |
| 349 | else: |
| 350 | rgb, disp, acc, extras = render(H, W, focal, chunk=args.chunk, c2w=pose_nerf[0,:3,:4], img_idx=img_idx, **render_kwargs_test) |
| 351 | rgb = rgb[None,...].permute(0,3,1,2) |
| 352 | |
| 353 | # feature metric module |
| 354 | feature_list, _ = inference_pose_regression(args, torch.cat([data, rgb]), device, feat_model, retFeature=True, isSingleStream=False, return_pose=False) |
| 355 | feature_target = feature_list[0] |
| 356 | feature_rgb = feature_list[1] |
| 357 | |
| 358 | ### Loss Design Here ### |
| 359 | # Compute RGB MSE Loss |
| 360 | photo_loss = rgb_loss(rgb, data, extras) |
| 361 | |
| 362 | # Compute Feature MSE Loss |
| 363 | indices = torch.tensor(args.feature_matching_lvl) |
| 364 | feature_rgb = torch.index_select(feature_rgb, 0, indices) |
| 365 | feature_target = torch.index_select(feature_target, 0, indices) |
| 366 | |
| 367 | feature_rgb = preprocess_features_for_loss(feature_rgb) |
| 368 | feature_target = preprocess_features_for_loss(feature_target) |
| 369 | |
| 370 | feat_loss = feature_loss(feature_rgb[0], feature_target[0], per_channel=args.per_channel) |
| 371 | |
| 372 | # Compute Combine Loss if needed |
| 373 | if args.combine_loss: |
| 374 | pose_loss = PoseLoss(args, pose_, pose, device) |
| 375 | loss = args.combine_loss_w[0] * pose_loss + args.combine_loss_w[1] * photo_loss + args.combine_loss_w[2] * feat_loss |
| 376 | |
| 377 | ### Loss Design End |
| 378 | loss.backward() |
| 379 | optimizer.step() |
no test coverage detected