目錄

• 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: 代碼閲讀-訓練參數

    # 命令行參數解析器
    parser = ArgumentParser(description="Training script parameters")
    # 模型相關參數
    lp = ModelParams(parser)
    op = OptimizationParams(parser)
    pp = PipelineParams(parser)
    parser.add_argument('--ip', type=str, default="127.0.0.1")
    parser.add_argument('--port', type=int, default=6009)
    parser.add_argument('--debug_from', type=int, default=-1)
    parser.add_argument('--detect_anomaly', action='store_true', default=False)
    parser.add_argument("--test_iterations", nargs="+", type=int, default=[7_000, 30_000])
    parser.add_argument("--save_iterations", nargs="+", type=int, default=[7_000, 30_000])
    parser.add_argument("--quiet", action="store_true")
    parser.add_argument('--disable_viewer', action='store_true', default=False)
    parser.add_argument("--checkpoint_iterations", nargs="+", type=int, default=[])
    parser.add_argument("--start_checkpoint", type=str, default = None)
    args = parser.parse_args(sys.argv[1:])

開始訓練

程序調用 training() 這個方法開始訓練

torch.autograd.set_detect_anomaly(args.detect_anomaly)
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() 這個方法中初始化訓練的代碼和對應的註釋説明

# 如果指定了 sparse_adam 加速器, 檢查是否已經安裝
if not SPARSE_ADAM_AVAILABLE and opt.optimizer_type == "sparse_adam":
    sys.exit(f"Trying to use sparse adam but it is not installed, please install the correct rasterizer using pip install [3dgs_accel].")
# first_iter用於記錄當前是第幾次迭代
first_iter = 0
# 創建本次訓練的輸出目錄和日誌記錄器, 每次執行訓練, 都會在 output 目錄下創建一個隨機目錄名
tb_writer = prepare_output_and_logger(dataset)
# 初始化 Gaussian 模型
gaussians = GaussianModel(dataset.sh_degree, opt.optimizer_type)
# 初始化訓練場景, 這裏會載入相機參數和稀疏點雲等數據
scene = Scene(dataset, gaussians)
# 初始化訓練參數
gaussians.training_setup(opt)
# 如果存在檢查點, 則載入
if checkpoint:
    (model_params, first_iter) = torch.load(checkpoint)
    gaussians.restore(model_params, opt)

# 設置背景顏色
bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]
background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")

# 初始化CUDA事件
iter_start = torch.cuda.Event(enable_timing = True)
iter_end = torch.cuda.Event(enable_timing = True)

# 是否使用 sparse adam 加速器
use_sparse_adam = opt.optimizer_type == "sparse_adam" and SPARSE_ADAM_AVAILABLE 
# Get depth L1 weight scheduling function
depth_l1_weight = get_expon_lr_func(opt.depth_l1_weight_init, opt.depth_l1_weight_final, max_steps=opt.iterations)

# Initialize viewpoint stack and indices
viewpoint_stack = scene.getTrainCameras().copy()
viewpoint_indices = list(range(len(viewpoint_stack)))
# Initialize exponential moving averages for logging
ema_loss_for_log = 0.0
ema_Ll1depth_for_log = 0.0

# 初始化進度條
progress_bar = tqdm(range(first_iter, opt.iterations), desc="Training progress")

迭代訓練

從大約73行開始, 進行迭代訓練

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

對外連工具展示渲染結果

    # 這部分處理網絡連接, 對外展示當前訓練的渲染結果
    if network_gui.conn == None:
        network_gui.try_connect()
    while network_gui.conn != None:
        try:
            net_image_bytes = None
            # Receive data from GUI
            custom_cam, do_training, pipe.convert_SHs_python, pipe.compute_cov3D_python, keep_alive, scaling_modifer = network_gui.receive()
            if custom_cam != None:
                # Render image for GUI
                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"]
                net_image_bytes = memoryview((torch.clamp(net_image, min=0, max=1.0) * 255).byte().permute(1, 2, 0).contiguous().cpu().numpy())
            # Send image to GUI
            network_gui.send(net_image_bytes, dataset.source_path)
            if do_training and ((iteration < int(opt.iterations)) or not keep_alive):
                break
        except Exception as e:
            network_gui.conn = None

更新學習率, 選中相機進行渲染

    # 記錄迭代的開始時間
    iter_start.record()

    # 更新學習率, 底下都是調用的 get_expon_lr_func(), 一個學習率調度函數, 根據訓練步數計算當前的學習率, 學習率從初始值指數衰減到最終值.
    gaussians.update_learning_rate(iteration)

    # 每1000次迭代, 球諧函數(SH, Spherical Harmonics)的階數加1, 直到設置的最大的階數, 默認最大為3, 
    # 每個3D高斯點需要存儲(階數 + 1)^2 個球諧係數, 3階時為16個係數, 每個係數有RGB 3個值所以一共48個值
    if iteration % 1000 == 0:
        gaussians.oneupSHdegree()

    # 當棧為空時, 複製一份訓練幀的相機位姿列表並創建對應的索引列表
    if not viewpoint_stack:
        viewpoint_stack = scene.getTrainCameras().copy()
        viewpoint_indices = list(range(len(viewpoint_stack)))
    # 從中隨機選取一個相機位姿
    rand_idx = randint(0, len(viewpoint_indices) - 1)
    # 從當前棧中彈出, 避免重複選取, 這樣最終會按隨機的順序遍歷完所有的相機位姿
    viewpoint_cam = viewpoint_stack.pop(rand_idx)
    vind = viewpoint_indices.pop(rand_idx)

    # 如果到了開啓debug的迭代次數, 開啓debug
    if (iteration - 1) == debug_from:
        pipe.debug = True

    # 如果設置了隨機背景, 創建隨機背景顏色張量
    bg = torch.rand((3), device="cuda") if opt.random_background else background

    # 用當前選中的相機視角, 渲染當前的場景
    render_pkg = render(viewpoint_cam, gaussians, pipe, bg, use_trained_exp=dataset.train_test_exp, separate_sh=SPARSE_ADAM_AVAILABLE)
    # 讀出渲染結果
    image, viewspace_point_tensor, visibility_filter, radii 
        = render_pkg["render"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["radii"]

    # 處理攝像機視角的alpha遮罩(透明度), 將alpha遮罩數據從CPU內存轉移到GPU顯存, 將當前圖像與alpha遮罩進行逐像素相乘, 
    # alpha值為1時保留原像素, alpha值為0時使像素完全透明
    if viewpoint_cam.alpha_mask is not None:
        alpha_mask = viewpoint_cam.alpha_mask.cuda()
        image *= alpha_mask

計算損失

    # 從viewpoint_cam對象中獲取原始圖像數據, 使用.cuda()方法將數據從CPU內存轉移到GPU顯存, 
    # 調用L1損失函數, 計算渲染結果與原圖gt_image之間的像素級絕對差平均值
    gt_image = viewpoint_cam.original_image.cuda()
    Ll1 = l1_loss(image, gt_image)
    # 計算兩個圖像之間的結構相似性指數(SSIM), 如果 fused_ssim 可用則使用 fused_ssim, 否則使用普通的ssim
    # Calculate SSIM using fused implementation if available
    if FUSED_SSIM_AVAILABLE:
        # 用unsqueeze(0)來增加一個維度，fused_ssim需要批量輸入
        ssim_value = fused_ssim(image.unsqueeze(0), gt_image.unsqueeze(0))
    else:
        ssim_value = ssim(image, gt_image)

    # 結合L1損失和SSIM損失計算混合損失, (1.0 - ssim_value) 將SSIM相似度轉換為損失值, 因為SSIM值越大損失越小
    loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim_value)

    # Depth regularization 深度正則化, 引入單目深度估計作為弱監督信號改善幾何一致性, 緩解漂浮物偽影, 增強遮擋區域的重建效果
    Ll1depth_pure = 0.0
    if depth_l1_weight(iteration) > 0 and viewpoint_cam.depth_reliable:
        # 從渲染結果中獲取逆向深度圖(1/depth)
        invDepth = render_pkg["depth"]
        # 獲取單目深度估計的逆向深度圖並轉移到GPU
        mono_invdepth = viewpoint_cam.invdepthmap.cuda()
        # 深度有效區域的掩碼(標記可靠區域)
        depth_mask = viewpoint_cam.depth_mask.cuda()

        # 計算帶掩碼的L1損失 = 絕對差(渲染深度 - 單目深度) * 掩碼 → 取均值
        Ll1depth_pure = torch.abs((invDepth  - mono_invdepth) * depth_mask).mean()
        # 應用動態權重係數(可能隨迭代次數衰減)
        Ll1depth = depth_l1_weight(iteration) * Ll1depth_pure 
        # 將加權後的深度損失加入總損失
        loss += Ll1depth
        # 將Tensor轉換為Python數值用於記錄
        Ll1depth = Ll1depth.item()
    else:
        Ll1depth = 0

反向計算梯度並優化

    # 執行反向傳播算法, 自動計算所有可訓練參數關於loss的梯度
    loss.backward()

    # 記錄迭代結束時間
    iter_end.record()
    # End iteration timing

    # torch.no_grad() 臨時關閉梯度計算的上下文管理器
    with torch.no_grad():
        # Progress bar
        ema_loss_for_log = 0.4 * loss.item() + 0.6 * ema_loss_for_log
        ema_Ll1depth_for_log = 0.4 * Ll1depth + 0.6 * ema_Ll1depth_for_log

        if iteration % 10 == 0:
            progress_bar.set_postfix({"Loss": f"{ema_loss_for_log:.{7}f}", "Depth Loss": f"{ema_Ll1depth_for_log:.{7}f}"})
            progress_bar.update(10)
        if iteration == opt.iterations:
            progress_bar.close()

        # 輸出日誌, 當迭代次數為 testing_iterations 時(默認為7000和30000), 會做一次整體評估, 間隔5取5個樣本, 取一部分相機視角計算L1和SSIM損失, iter_start.elapsed_time(iter_end) 計算耗時
        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)
        # 當迭代次數為 saving_iterations(默認為7000和30000)時,保存
        if (iteration in saving_iterations):
            print("\n[ITER {}] Saving Gaussians".format(iteration))
            # 裏面會調用 gaussians.save_ply() 保存ply文件
            scene.save(iteration)

        # 當迭代次數小於緻密化結束的右邊界時
        if iteration < opt.densify_until_iter:
            # 可見性半徑更新, 記錄每個高斯點在所有視角下的最大可見半徑, 用於後續剪枝判斷. visibility_filter過濾出當前視角可見的高斯點
            # Keep track of max radii in image-space for pruning
            gaussians.max_radii2D[visibility_filter] = torch.max(gaussians.max_radii2D[visibility_filter], radii[visibility_filter])
            # 累積視空間位置梯度統計量, 用於後續判斷哪些高斯點需要分裂(高梯度區域)或克隆(高位置變化區域)
            gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter)

            # 當迭代次數大於緻密化開始的左邊界, 並且滿足緻密化間隔時, 進行緻密化與修剪處理
            if iteration > opt.densify_from_iter and iteration % opt.densification_interval == 0:
                # 如果迭代次數小於不透明度重置間隔(3000)則返回20作為2D尺寸限制, 否則不限制
                size_threshold = 20 if iteration > opt.opacity_reset_interval else None
                # 緻密化與修剪
                gaussians.densify_and_prune(opt.densify_grad_threshold, 0.005, scene.cameras_extent, size_threshold, radii)
            
            # 定期(默認3000一次)重置不透明度, 恢復被錯誤剪枝的高斯點, 調整新生成高斯的可見性, 適配白背景場景的特殊初始化
            if iteration % opt.opacity_reset_interval == 0 or (dataset.white_background and iteration == opt.densify_from_iter):
                gaussians.reset_opacity()

        # Optimizer階段, 反向優化模型參數
        if iteration < opt.iterations:
            gaussians.exposure_optimizer.step()
            gaussians.exposure_optimizer.zero_grad(set_to_none = True)
            if use_sparse_adam:
                visible = radii > 0
                gaussians.optimizer.step(visible, radii.shape[0])
                gaussians.optimizer.zero_grad(set_to_none = True)
            else:
                gaussians.optimizer.step()
                gaussians.optimizer.zero_grad(set_to_none = True)

        # 到達預設的checkpoint, 默認為7000和30000, 保存當前的訓練進度
        if (iteration in checkpoint_iterations):
            print("\n[ITER {}] Saving Checkpoint".format(iteration))
            torch.save((gaussians.capture(), iteration), scene.model_path + "/chkpnt" + str(iteration) + ".pth")

如果有錯誤請留言指出