3D Gaussian splatting 05: 代码阅读-训练整体流程

目录

• 3D Gaussian splatting 01: 环境搭建
  
• 3D Gaussian splatting 02: 快速评估
  
• 3D Gaussian splatting 03: 用户数据训练和结果查看
  
• 3D Gaussian splatting 04: 代码阅读-提取相机位姿和稀疏点云
  
• 3D Gaussian splatting 05: 代码阅读-训练整体流程
  
• 3D Gaussian splatting 06: 代码阅读-训练参数
  
• 3D Gaussian splatting 07: 代码阅读-训练载入数据和保存结果
  
• 3D Gaussian splatting 08: 代码阅读-渲染
  

训练整体流程

程序入参

训练程序入参除了训练过程参数, 另外设置了ModelParams, OptimizationParams, PipelineParams三个参数组, 分别控制数据加载、渲染计算和优化训练环节, 详细的说明查看下一节 06: 代码阅读-训练参数

1.     # 命令行参数解析器
2.     parser = ArgumentParser(description="Training script parameters")
3.     # 模型相关参数
4.     lp = ModelParams(parser)
5.     op = OptimizationParams(parser)
6.     pp = PipelineParams(parser)
7.     parser.add_argument('--ip', type=str, default="127.0.0.1")
8.     parser.add_argument('--port', type=int, default=6009)
9.     parser.add_argument('--debug_from', type=int, default=-1)
10.     parser.add_argument('--detect_anomaly', action='store_true', default=False)
11.     parser.add_argument("--test_iterations", nargs="+", type=int, default=[7_000, 30_000])
12.     parser.add_argument("--save_iterations", nargs="+", type=int, default=[7_000, 30_000])
13.     parser.add_argument("--quiet", action="store_true")
14.     parser.add_argument('--disable_viewer', action='store_true', default=False)
15.     parser.add_argument("--checkpoint_iterations", nargs="+", type=int, default=[])
16.     parser.add_argument("--start_checkpoint", type=str, default = None)
17.     args = parser.parse_args(sys.argv[1:])

复制代码

开始训练

程序调用 training() 这个方法开始训练

1. torch.autograd.set_detect_anomaly(args.detect_anomaly)
2. training(lp.extract(args), op.extract(args), pp.extract(args), args.test_iterations, args.save_iterations, args.checkpoint_iterations, args.start_checkpoint, args.debug_from)

复制代码

初始化

以下是 training() 这个方法中初始化训练的代码和对应的注释说明

1. # 如果指定了 sparse_adam 加速器, 检查是否已经安装
2. if not SPARSE_ADAM_AVAILABLE and opt.optimizer_type == "sparse_adam":
3.     sys.exit(f"Trying to use sparse adam but it is not installed, please install the correct rasterizer using pip install [3dgs_accel].")
4. # first_iter用于记录当前是第几次迭代
5. first_iter = 0
6. # 创建本次训练的输出目录和日志记录器, 每次执行训练, 都会在 output 目录下创建一个随机目录名
7. tb_writer = prepare_output_and_logger(dataset)
8. # 初始化 Gaussian 模型
9. gaussians = GaussianModel(dataset.sh_degree, opt.optimizer_type)
10. # 初始化训练场景, 这里会载入相机参数和稀疏点云等数据
11. scene = Scene(dataset, gaussians)
12. # 初始化训练参数
13. gaussians.training_setup(opt)
14. # 如果存在检查点, 则载入
15. if checkpoint:
16.     (model_params, first_iter) = torch.load(checkpoint)
17.     gaussians.restore(model_params, opt)
19. # 设置背景颜色
20. bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]
21. background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
23. # 初始化CUDA事件
24. iter_start = torch.cuda.Event(enable_timing = True)
25. iter_end = torch.cuda.Event(enable_timing = True)
27. # 是否使用 sparse adam 加速器
28. use_sparse_adam = opt.optimizer_type == "sparse_adam" and SPARSE_ADAM_AVAILABLE
29. # Get depth L1 weight scheduling function
30. depth_l1_weight = get_expon_lr_func(opt.depth_l1_weight_init, opt.depth_l1_weight_final, max_steps=opt.iterations)
32. # Initialize viewpoint stack and indices
33. viewpoint_stack = scene.getTrainCameras().copy()
34. viewpoint_indices = list(range(len(viewpoint_stack)))
35. # Initialize exponential moving averages for logging
36. ema_loss_for_log = 0.0
37. ema_Ll1depth_for_log = 0.0
39. # 初始化进度条
40. progress_bar = tqdm(range(first_iter, opt.iterations), desc="Training progress")

复制代码

迭代训练

从大约73行开始, 进行迭代训练

1. first_iter += 1
2. for iteration in range(first_iter, opt.iterations + 1):

复制代码

对外连工具展示渲染结果

1.     # 这部分处理网络连接, 对外展示当前训练的渲染结果
2.     if network_gui.conn == None:
3.         network_gui.try_connect()
4.     while network_gui.conn != None:
5.         try:
6.             net_image_bytes = None
7.             # Receive data from GUI
8.             custom_cam, do_training, pipe.convert_SHs_python, pipe.compute_cov3D_python, keep_alive, scaling_modifer = network_gui.receive()
9.             if custom_cam != None:
10.                 # Render image for GUI
11.                 net_image = render(custom_cam, gaussians, pipe, background, scaling_modifier=scaling_modifer, use_trained_exp=dataset.train_test_exp, separate_sh=SPARSE_ADAM_AVAILABLE)["render"]
12.                 net_image_bytes = memoryview((torch.clamp(net_image, min=0, max=1.0) * 255).byte().permute(1, 2, 0).contiguous().cpu().numpy())
13.             # Send image to GUI
14.             network_gui.send(net_image_bytes, dataset.source_path)
15.             if do_training and ((iteration < int(opt.iterations)) or not keep_alive):
16.                 break
17.         except Exception as e:
18.             network_gui.conn = None

复制代码

更新学习率, 选中相机进行渲染

1.     # 记录迭代的开始时间
2.     iter_start.record()
4.     # 更新学习率, 底下都是调用的 get_expon_lr_func(), 一个学习率调度函数, 根据训练步数计算当前的学习率, 学习率从初始值指数衰减到最终值.
5.     gaussians.update_learning_rate(iteration)
7.     # 每1000次迭代, 球谐函数(SH, Spherical Harmonics)的阶数加1, 直到设置的最大的阶数, 默认最大为3,
8.     # 每个3D高斯点需要存储(阶数 + 1)^2 个球谐系数, 3阶时为16个系数, 每个系数有RGB 3个值所以一共48个值
9.     if iteration % 1000 == 0:
10.         gaussians.oneupSHdegree()
12.     # 当栈为空时, 复制一份训练帧的相机位姿列表并创建对应的索引列表
13.     if not viewpoint_stack:
14.         viewpoint_stack = scene.getTrainCameras().copy()
15.         viewpoint_indices = list(range(len(viewpoint_stack)))
16.     # 从中随机选取一个相机位姿
17.     rand_idx = randint(0, len(viewpoint_indices) - 1)
18.     # 从当前栈中弹出, 避免重复选取, 这样最终会按随机的顺序遍历完所有的相机位姿
19.     viewpoint_cam = viewpoint_stack.pop(rand_idx)
20.     vind = viewpoint_indices.pop(rand_idx)
22.     # 如果到了开启debug的迭代次数, 开启debug
23.     if (iteration - 1) == debug_from:
24.         pipe.debug = True
26.     # 如果设置了随机背景, 创建随机背景颜色张量
27.     bg = torch.rand((3), device="cuda") if opt.random_background else background
29.     # 用当前选中的相机视角, 渲染当前的场景
30.     render_pkg = render(viewpoint_cam, gaussians, pipe, bg, use_trained_exp=dataset.train_test_exp, separate_sh=SPARSE_ADAM_AVAILABLE)
31.     # 读出渲染结果
32.     image, viewspace_point_tensor, visibility_filter, radii
33.         = render_pkg["render"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["radii"]
35.     # 处理摄像机视角的alpha遮罩(透明度), 将alpha遮罩数据从CPU内存转移到GPU显存, 将当前图像与alpha遮罩进行逐像素相乘,
36.     # alpha值为1时保留原像素, alpha值为0时使像素完全透明
37.     if viewpoint_cam.alpha_mask is not None:
38.         alpha_mask = viewpoint_cam.alpha_mask.cuda()
39.         image *= alpha_mask

复制代码

计算损失

1.     # 从viewpoint_cam对象中获取原始图像数据, 使用.cuda()方法将数据从CPU内存转移到GPU显存,
2.     # 调用L1损失函数, 计算渲染结果与原图gt_image之间的像素级绝对差平均值
3.     gt_image = viewpoint_cam.original_image.cuda()
4.     Ll1 = l1_loss(image, gt_image)
5.     # 计算两个图像之间的结构相似性指数(SSIM), 如果 fused_ssim 可用则使用 fused_ssim, 否则使用普通的ssim
6.     # Calculate SSIM using fused implementation if available
7.     if FUSED_SSIM_AVAILABLE:
8.         # 用unsqueeze(0)来增加一个维度，fused_ssim需要批量输入
9.         ssim_value = fused_ssim(image.unsqueeze(0), gt_image.unsqueeze(0))
10.     else:
11.         ssim_value = ssim(image, gt_image)
13.     # 结合L1损失和SSIM损失计算混合损失, (1.0 - ssim_value) 将SSIM相似度转换为损失值, 因为SSIM值越大损失越小
14.     loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim_value)
16.     # Depth regularization 深度正则化, 引入单目深度估计作为弱监督信号改善几何一致性, 缓解漂浮物伪影, 增强遮挡区域的重建效果
17.     Ll1depth_pure = 0.0
18.     if depth_l1_weight(iteration) > 0 and viewpoint_cam.depth_reliable:
19.         # 从渲染结果中获取逆向深度图(1/depth)
20.         invDepth = render_pkg["depth"]
21.         # 获取单目深度估计的逆向深度图并转移到GPU
22.         mono_invdepth = viewpoint_cam.invdepthmap.cuda()
23.         # 深度有效区域的掩码(标记可靠区域)
24.         depth_mask = viewpoint_cam.depth_mask.cuda()
26.         # 计算带掩码的L1损失 = 绝对差(渲染深度 - 单目深度) * 掩码 → 取均值
27.         Ll1depth_pure = torch.abs((invDepth  - mono_invdepth) * depth_mask).mean()
28.         # 应用动态权重系数(可能随迭代次数衰减)
29.         Ll1depth = depth_l1_weight(iteration) * Ll1depth_pure
30.         # 将加权后的深度损失加入总损失
31.         loss += Ll1depth
32.         # 将Tensor转换为Python数值用于记录
33.         Ll1depth = Ll1depth.item()
34.     else:
35.         Ll1depth = 0

复制代码

反向计算梯度并优化

1.     # 执行反向传播算法, 自动计算所有可训练参数关于loss的梯度
2.     loss.backward()
4.     # 记录迭代结束时间
5.     iter_end.record()
6.     # End iteration timing
8.     # torch.no_grad() 临时关闭梯度计算的上下文管理器
9.     with torch.no_grad():
10.         # Progress bar
11.         ema_loss_for_log = 0.4 * loss.item() + 0.6 * ema_loss_for_log
12.         ema_Ll1depth_for_log = 0.4 * Ll1depth + 0.6 * ema_Ll1depth_for_log
14.         if iteration % 10 == 0:
15.             progress_bar.set_postfix({"Loss": f"{ema_loss_for_log:.{7}f}", "Depth Loss": f"{ema_Ll1depth_for_log:.{7}f}"})
16.             progress_bar.update(10)
17.         if iteration == opt.iterations:
18.             progress_bar.close()
20.         # 输出日志, 当迭代次数为 testing_iterations 时(默认为7000和30000), 会做一次整体评估, 间隔5取5个样本, 取一部分相机视角计算L1和SSIM损失, iter_start.elapsed_time(iter_end) 计算耗时
21.         training_report(tb_writer, iteration, Ll1, loss, l1_loss, iter_start.elapsed_time(iter_end), testing_iterations, scene, render, (pipe, background, 1., SPARSE_ADAM_AVAILABLE, None, dataset.train_test_exp), dataset.train_test_exp)
22.         # 当迭代次数为 saving_iterations(默认为7000和30000)时,保存
23.         if (iteration in saving_iterations):
24.             print("\n[ITER {}] Saving Gaussians".format(iteration))
25.             # 里面会调用 gaussians.save_ply() 保存ply文件
26.             scene.save(iteration)
28.         # 当迭代次数小于致密化结束的右边界时
29.         if iteration < opt.densify_until_iter:
30.             # 可见性半径更新, 记录每个高斯点在所有视角下的最大可见半径, 用于后续剪枝判断. visibility_filter过滤出当前视角可见的高斯点
31.             # Keep track of max radii in image-space for pruning
32.             gaussians.max_radii2D[visibility_filter] = torch.max(gaussians.max_radii2D[visibility_filter], radii[visibility_filter])
33.             # 累积视空间位置梯度统计量, 用于后续判断哪些高斯点需要分裂(高梯度区域)或克隆(高位置变化区域)
34.             gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter)
36.             # 当迭代次数大于致密化开始的左边界, 并且满足致密化间隔时, 进行致密化与修剪处理
37.             if iteration > opt.densify_from_iter and iteration % opt.densification_interval == 0:
38.                 # 如果迭代次数小于不透明度重置间隔(3000)则返回20作为2D尺寸限制, 否则不限制
39.                 size_threshold = 20 if iteration > opt.opacity_reset_interval else None
40.                 # 致密化与修剪
41.                 gaussians.densify_and_prune(opt.densify_grad_threshold, 0.005, scene.cameras_extent, size_threshold, radii)
42.             
43.             # 定期(默认3000一次)重置不透明度, 恢复被错误剪枝的高斯点, 调整新生成高斯的可见性, 适配白背景场景的特殊初始化
44.             if iteration % opt.opacity_reset_interval == 0 or (dataset.white_background and iteration == opt.densify_from_iter):
45.                 gaussians.reset_opacity()
47.         # Optimizer阶段, 反向优化模型参数
48.         if iteration < opt.iterations:
49.             gaussians.exposure_optimizer.step()
50.             gaussians.exposure_optimizer.zero_grad(set_to_none = True)
51.             if use_sparse_adam:
52.                 visible = radii > 0
53.                 gaussians.optimizer.step(visible, radii.shape[0])
54.                 gaussians.optimizer.zero_grad(set_to_none = True)
55.             else:
56.                 gaussians.optimizer.step()
57.                 gaussians.optimizer.zero_grad(set_to_none = True)
59.         # 到达预设的checkpoint, 默认为7000和30000, 保存当前的训练进度
60.         if (iteration in checkpoint_iterations):
61.             print("\n[ITER {}] Saving Checkpoint".format(iteration))
62.             torch.save((gaussians.capture(), iteration), scene.model_path + "/chkpnt" + str(iteration) + ".pth")

复制代码

如果有错误请留言指出

来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！