(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoint_iterations, checkpoint, debug_from)
| 86 | |
| 87 | |
| 88 | def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoint_iterations, checkpoint, debug_from): |
| 89 | |
| 90 | first_iter = 0 |
| 91 | tb_writer = prepare_output_and_logger(dataset) |
| 92 | gaussians = GaussianModel(dataset.sh_degree) |
| 93 | |
| 94 | # per-point-optimizer |
| 95 | confidence_path = os.path.join(dataset.source_path, f"sparse_{dataset.n_views}/0", "confidence_dsp.npy") |
| 96 | confidence_lr = load_and_prepare_confidence(confidence_path, device='cuda', scale=(1, 100)) |
| 97 | scene = Scene(dataset, gaussians) |
| 98 | |
| 99 | if opt.pp_optimizer: |
| 100 | gaussians.training_setup_pp(opt, confidence_lr) |
| 101 | else: |
| 102 | gaussians.training_setup(opt) |
| 103 | if checkpoint: |
| 104 | (model_params, first_iter) = torch.load(checkpoint) |
| 105 | gaussians.restore(model_params, opt) |
| 106 | |
| 107 | train_cams_init = scene.getTrainCameras().copy() |
| 108 | for save_iter in saving_iterations: |
| 109 | os.makedirs(scene.model_path + f'/pose/ours_{save_iter}', exist_ok=True) |
| 110 | save_pose(scene.model_path + f'/pose/ours_{save_iter}/pose_org.npy', gaussians.P, train_cams_init) |
| 111 | bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0] |
| 112 | background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") |
| 113 | |
| 114 | iter_start = torch.cuda.Event(enable_timing = True) |
| 115 | iter_end = torch.cuda.Event(enable_timing = True) |
| 116 | |
| 117 | viewpoint_stack = scene.getTrainCameras().copy() |
| 118 | viewpoint_indices = list(range(len(viewpoint_stack))) |
| 119 | ema_loss_for_log = 0.0 |
| 120 | |
| 121 | progress_bar = tqdm(range(first_iter, opt.iterations), desc="Training progress") |
| 122 | first_iter += 1 |
| 123 | start = time() |
| 124 | for iteration in range(first_iter, opt.iterations + 1): |
| 125 | # if network_gui.conn == None: |
| 126 | # network_gui.try_connect() |
| 127 | # while network_gui.conn != None: |
| 128 | # try: |
| 129 | # net_image_bytes = None |
| 130 | # custom_cam, do_training, pipe.convert_SHs_python, pipe.compute_cov3D_python, keep_alive, scaling_modifer = network_gui.receive() |
| 131 | # if custom_cam != None: |
| 132 | # net_image = render(custom_cam, gaussians, pipe, background, scaling_modifer)["render"] |
| 133 | # net_image_bytes = memoryview((torch.clamp(net_image, min=0, max=1.0) * 255).byte().permute(1, 2, 0).contiguous().cpu().numpy()) |
| 134 | # network_gui.send(net_image_bytes, dataset.source_path) |
| 135 | # if do_training and ((iteration < int(opt.iterations)) or not keep_alive): |
| 136 | # break |
| 137 | # except Exception as e: |
| 138 | # network_gui.conn = None |
| 139 | |
| 140 | iter_start.record() |
| 141 | |
| 142 | gaussians.update_learning_rate(iteration) |
| 143 | |
| 144 | if opt.optim_pose==False: |
| 145 | gaussians.P.requires_grad_(False) |
no test coverage detected