(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoint, debug_from,
gaussian_dim, time_duration, num_pts, num_pts_ratio, rot_4d, force_sh_3d, batch_size)
| 35 | TENSORBOARD_FOUND = False |
| 36 | |
| 37 | def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoint, debug_from, |
| 38 | gaussian_dim, time_duration, num_pts, num_pts_ratio, rot_4d, force_sh_3d, batch_size): |
| 39 | |
| 40 | if dataset.frame_ratio > 1: |
| 41 | time_duration = [time_duration[0] / dataset.frame_ratio, time_duration[1] / dataset.frame_ratio] |
| 42 | |
| 43 | first_iter = 0 |
| 44 | tb_writer = prepare_output_and_logger(dataset) |
| 45 | gaussians = GaussianModel(dataset.sh_degree, gaussian_dim=gaussian_dim, time_duration=time_duration, rot_4d=rot_4d, force_sh_3d=force_sh_3d, sh_degree_t=2 if pipe.eval_shfs_4d else 0) |
| 46 | scene = Scene(dataset, gaussians, num_pts=num_pts, num_pts_ratio=num_pts_ratio, time_duration=time_duration) |
| 47 | gaussians.training_setup(opt) |
| 48 | |
| 49 | if checkpoint: |
| 50 | (model_params, first_iter) = torch.load(checkpoint) |
| 51 | gaussians.restore(model_params, opt) |
| 52 | |
| 53 | bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0] |
| 54 | background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") |
| 55 | |
| 56 | iter_start = torch.cuda.Event(enable_timing = True) |
| 57 | iter_end = torch.cuda.Event(enable_timing = True) |
| 58 | |
| 59 | best_psnr = 0.0 |
| 60 | ema_loss_for_log = 0.0 |
| 61 | ema_l1loss_for_log = 0.0 |
| 62 | ema_ssimloss_for_log = 0.0 |
| 63 | lambda_all = [key for key in opt.__dict__.keys() if key.startswith('lambda') and key!='lambda_dssim'] |
| 64 | for lambda_name in lambda_all: |
| 65 | vars()[f"ema_{lambda_name.replace('lambda_','')}_for_log"] = 0.0 |
| 66 | |
| 67 | progress_bar = tqdm(range(first_iter, opt.iterations), desc="Training progress") |
| 68 | first_iter += 1 |
| 69 | |
| 70 | if pipe.env_map_res: |
| 71 | env_map = nn.Parameter(torch.zeros((3,pipe.env_map_res, pipe.env_map_res),dtype=torch.float, device="cuda").requires_grad_(True)) |
| 72 | env_map_optimizer = torch.optim.Adam([env_map], lr=opt.feature_lr, eps=1e-15) |
| 73 | else: |
| 74 | env_map = None |
| 75 | |
| 76 | gaussians.env_map = env_map |
| 77 | |
| 78 | training_dataset = scene.getTrainCameras() |
| 79 | # import pdb |
| 80 | # pdb.set_trace() |
| 81 | training_dataloader = DataLoader(training_dataset, batch_size=batch_size, shuffle=True, num_workers=12 if dataset.dataloader else 0, collate_fn=lambda x: x, drop_last=True) |
| 82 | |
| 83 | iteration = first_iter |
| 84 | while iteration < opt.iterations + 1: |
| 85 | for batch_data in training_dataloader: |
| 86 | iteration += 1 |
| 87 | if iteration > opt.iterations: |
| 88 | break |
| 89 | |
| 90 | iter_start.record() |
| 91 | gaussians.update_learning_rate(iteration) |
| 92 | |
| 93 | # Every 1000 its we increase the levels of SH up to a maximum degree |
| 94 | if iteration % opt.sh_increase_interval == 0: |
no test coverage detected