(args)
| 38 | |
| 39 | |
| 40 | def train(args): |
| 41 | device = torch.device(args.render_device) |
| 42 | dataset_name = args.scene_name |
| 43 | log_file_path = os.path.join(args.model_path, "logs", |
| 44 | f"({time.strftime('%Y-%m-%d_%H-%M-%S')})_iteration({args.iterations})_({dataset_name}).log") |
| 45 | os.makedirs(os.path.dirname(log_file_path), exist_ok=True) |
| 46 | logging.basicConfig(filename=log_file_path, level=logging.INFO, |
| 47 | format='%(asctime)s - %(levelname)s - %(message)s') |
| 48 | logging.info(f"Experiment Configuration: {args}") |
| 49 | logging.info(f"Model initialization and Data reading ...") |
| 50 | # save args |
| 51 | with open(os.path.join(args.model_path, 'cfg_arg.pkl'), 'wb') as file: |
| 52 | pickle.dump(args, file) |
| 53 | first_iter = 0 |
| 54 | |
| 55 | tb_writer = prepare_output_and_logger(args) |
| 56 | if args.depth_is_inverted: |
| 57 | from models.gs.gaussian_model_inv_depth import GaussianModelInvDepth |
| 58 | gaussians = GaussianModelInvDepth(args) |
| 59 | else: |
| 60 | gaussians = GaussianModel(args) |
| 61 | scene = Scene(args, gaussians, shuffle=False) |
| 62 | gaussians.training_setup(args) |
| 63 | |
| 64 | if args.deblur: |
| 65 | blur_blend_embedding = torch.nn.Embedding( |
| 66 | len(scene.train_cameras), args.blur_sample_num, device=device) |
| 67 | blur_blend_embedding.weight = torch.nn.Parameter(torch.ones( |
| 68 | len(scene.train_cameras), args.blur_sample_num, device=device)) |
| 69 | optimizer = torch.optim.Adam([ |
| 70 | {'params': blur_blend_embedding.parameters(), |
| 71 | 'lr': 1e-3, "name": "blur blend parameters"}, |
| 72 | ], lr=0.0, eps=1e-15) |
| 73 | scheduler = torch.optim.lr_scheduler.ExponentialLR( |
| 74 | optimizer, gamma=(1e-6 / 1e-3) ** (1. / args.iterations)) |
| 75 | else: |
| 76 | args.blur_sample_num = 1 |
| 77 | |
| 78 | render_temp_path = os.path.join(args.model_path, "train_temp_rendering") |
| 79 | gt_temp_path = os.path.join(args.model_path, "train_temp_gt") |
| 80 | if os.path.exists(render_temp_path): |
| 81 | shutil.rmtree(render_temp_path) |
| 82 | if os.path.exists(gt_temp_path): |
| 83 | shutil.rmtree(gt_temp_path) |
| 84 | os.makedirs(render_temp_path, exist_ok=True) |
| 85 | os.makedirs(gt_temp_path, exist_ok=True) |
| 86 | |
| 87 | if args.use_features_mask: |
| 88 | render_temp_mask_path = os.path.join(args.model_path, "train_mask_temp_rendering") |
| 89 | if os.path.exists(render_temp_mask_path): |
| 90 | shutil.rmtree(render_temp_mask_path) |
| 91 | os.makedirs(render_temp_mask_path, exist_ok=True) |
| 92 | |
| 93 | bg_color = [1, 1, 1] if args.white_background else [0, 0, 0] |
| 94 | background = torch.tensor(bg_color, dtype=torch.float, device=device) |
| 95 | |
| 96 | iter_start = torch.cuda.Event(enable_timing=True) |
| 97 | iter_end = torch.cuda.Event(enable_timing=True) |
no test coverage detected