Perform 1 epoch of training with batch
(args, data_loaders, model, feat_model, hwf, optimizer, half_res, device, world_setup_dict, **render_kwargs_test)
| 390 | return iter_loss, iter_psnr |
| 391 | |
| 392 | def train_on_epoch(args, data_loaders, model, feat_model, hwf, optimizer, half_res, device, world_setup_dict, **render_kwargs_test): |
| 393 | ''' Perform 1 epoch of training with batch ''' |
| 394 | model.train() |
| 395 | model = freeze_bn_layer_train(model) |
| 396 | |
| 397 | # Prepare dataloaders for PoseNet, each batch contains (image, pose) |
| 398 | train_dl, val_dl, test_dl = data_loaders |
| 399 | total_loss = [] |
| 400 | total_psnr = [] |
| 401 | |
| 402 | #### Core optimization loop ##### |
| 403 | for data, pose, img_idx in train_dl: |
| 404 | loss, psnr = train_on_batch(args, data, model, feat_model, pose, img_idx, hwf, optimizer, half_res, device, world_setup_dict, **render_kwargs_test) |
| 405 | |
| 406 | total_loss.append(loss.item()) |
| 407 | total_psnr.append(psnr.item()) |
| 408 | total_loss_mean = np.mean(total_loss) |
| 409 | total_psnr_mean = np.mean(total_psnr) |
| 410 | return total_loss_mean, total_psnr_mean |
| 411 | |
| 412 | def train_feature_matching(args, model, feat_model, optimizer, i_split, hwf, near, far, device, early_stopping, images=None, poses_train=None, train_dl=None, val_dl=None, test_dl=None): |
| 413 | ''' finetune pretrained PoseNet using NeRF ''' |
no test coverage detected