3D Gaussian Splatting代码中的train和render两个文件代码解读

现在来聊一聊训练和渲染是如何进行的

training

train.py
line 31
def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoint_iterations, checkpoint, debug_from):
    # 初始化第一次迭代的索引为0
    first_iter = 0
    
    # 准备输出和日志记录器
    tb_writer = prepare_output_and_logger(dataset)
    
    # 初始化高斯模型，参数为数据集的球谐函数（SH）级别
    gaussians = GaussianModel(dataset.sh_degree)
    
    # 创建场景对象，包含数据集和高斯模型
    scene = Scene(dataset, gaussians)
    
    # 设置高斯模型的训练配置
    gaussians.training_setup(opt)
    
    # 加载检查点（如果有），恢复模型参数和设置起始迭代次数
    if checkpoint:
        (model_params, first_iter) = torch.load(checkpoint)
        gaussians.restore(model_params, opt)

    # 设置背景颜色，如果数据集背景为白色，则设置为白色（[1, 1, 1]），否则为黑色（[0, 0, 0]）
    bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]
    
    # 将背景颜色转换为CUDA张量，以便在GPU上使用
    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)

    # 初始化视点堆栈为空
    viewpoint_stack = None
    
    # 用于记录指数移动平均损失的变量，初始值为0.0
    ema_loss_for_log = 0.0
    
    # 创建进度条，用于显示训练进度，从起始迭代数到总迭代数
    progress_bar = tqdm(range(first_iter, opt.iterations), desc="Training progress")
    
    # 增加起始迭代数，以便从下一次迭代开始
    first_iter += 1

    for iteration in range(first_iter, opt.iterations + 1):
        # 尝试连接网络GUI，如果当前没有连接
        if network_gui.conn == None:
            network_gui.try_connect()

        # 如果已经连接网络GUI，处理接收和发送数据
        while network_gui.conn != None:
            try:
                # 初始化网络图像字节为None
                net_image_bytes = None

                # 从网络GUI接收数据
                custom_cam, do_training, pipe.convert_SHs_python, pipe.compute_cov3D_python, keep_alive, scaling_modifer = network_gui.receive()

                # 如果接收到自定义相机数据，则进行渲染
                if custom_cam != None:
                    # 使用自定义相机数据、当前的高斯模型、管道和背景颜色进行渲染
                    net_image = render(custom_cam, gaussians, pipe, background, scaling_modifer)["render"]

                    # 将渲染结果转为字节格式，并转换为内存视图
                    net_image_bytes = memoryview(
                        (torch.clamp(net_image, min=0, max=1.0) * 255).byte()
                        .permute(1, 2, 0).contiguous().cpu().numpy()
                    )

                # 发送渲染结果到网络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()


        # 更新学习率
        gaussians.update_learning_rate(iteration)

        # 每1000次迭代增加一次SH级别，直到达到最大度
        if iteration % 1000 == 0:
            gaussians.oneupSHdegree()

        # 随机选择一个相机视角
        if not viewpoint_stack:
            viewpoint_stack = scene.getTrainCameras().copy()
        # 从相机视角堆栈中随机弹出一个相机视角
        viewpoint_cam = viewpoint_stack.pop(randint(0, len(viewpoint_stack) - 1))

        # 渲染
        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)
        # 提取渲染结果、视点空间点张量、可见性过滤器和半径
        image, viewspace_point_tensor, visibility_filter, radii = render_pkg["render"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["radii"]

        # 计算损失
        gt_image = viewpoint_cam.original_image.cuda()  # 获取地面真实图像
        Ll1 = l1_loss(image, gt_image)  # 计算L1损失
        # 计算总损失，结合L1损失和结构相似性损失（SSIM）
        loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim(image, gt_image))
        loss.backward()  # 反向传播计算梯度

        # 记录当前迭代的结束时间，用于计算每次迭代的持续时间
        iter_end.record()

        # 在不需要计算梯度的上下文中进行操作
        with torch.no_grad():
            # 更新进度条和日志
            ema_loss_for_log = 0.4 * loss.item() + 0.6 * ema_loss_for_log  # 更新指数移动平均损失
            if iteration % 10 == 0:
                # 每10次迭代更新一次进度条
                progress_bar.set_postfix({"Loss": f"{ema_loss_for_log:.{7}f}"})
                progress_bar.update(10)
            if iteration == opt.iterations:
                progress_bar.close()

            # 记录训练报告并保存
            training_report(tb_writer, iteration, Ll1, loss, l1_loss, iter_start.elapsed_time(iter_end), testing_iterations, scene, render, (pipe, background))
            if iteration in saving_iterations:
                # 在指定的迭代次数保存高斯模型
                print("\n[ITER {}] Saving Gaussians".format(iteration))
                scene.save(iteration)

            # 密集化操作
            if iteration < opt.densify_until_iter:
                # 跟踪图像空间中的最大半径，用于修剪
                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:
                    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)

                # 在指定的间隔内或满足特定条件时重置不透明度
                if iteration % opt.opacity_reset_interval == 0 or (dataset.white_background and iteration == opt.densify_from_iter):
                    gaussians.reset_opacity()

            # 优化器步骤
            if iteration < opt.iterations:
                gaussians.optimizer.step()  # 更新模型参数
                gaussians.optimizer.zero_grad(set_to_none=True)  # 清空梯度

            # 保存检查点
            if iteration in checkpoint_iterations:
                print("\n[ITER {}] Saving Checkpoint".format(iteration))
                torch.save((gaussians.capture(), iteration), scene.model_path + "/chkpnt" + str(iteration) + ".pth")

render

现在是渲染的这个文件进行方式,首先是主文件里单张图片的渲染和整个数据集的渲染方法：

render.py
line 24
# 渲染一组视角并保存渲染结果和对应的真实图像
def render_set(model_path, name, iteration, views, gaussians, pipeline, background):
    # 定义渲染结果和真实图像的保存路径
    render_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders")
    gts_path = os.path.join(model_path, name, "ours_{}".format(iteration), "gt")

    # 创建保存路径，如果路径不存在
    makedirs(render_path, exist_ok=True)
    makedirs(gts_path, exist_ok=True)

    # 遍历每个视角进行渲染
    for idx, view in enumerate(tqdm(views, desc="Rendering progress")):
        # 渲染图像
        rendering = render(view, gaussians, pipeline, background)["render"]
        # 获取对应的真实图像
        gt = view.original_image[0:3, :, :]
        # 保存渲染结果和真实图像到指定路径
        torchvision.utils.save_image(rendering, os.path.join(render_path, '{0:05d}'.format(idx) + ".png"))
        torchvision.utils.save_image(gt, os.path.join(gts_path, '{0:05d}'.format(idx) + ".png"))

# 渲染训练集和测试集的图像，并保存结果
def render_sets(dataset: ModelParams, iteration: int, pipeline: PipelineParams, skip_train: bool, skip_test: bool):
    with torch.no_grad():
        # 初始化高斯模型和场景
        gaussians = GaussianModel(dataset.sh_degree)
        scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False)

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

        # 如果不跳过训练集渲染，则渲染训练集的图像
        if not skip_train:
            render_set(dataset.model_path, "train", scene.loaded_iter, scene.getTrainCameras(), gaussians, pipeline, background)

        # 如果不跳过测试集渲染，则渲染测试集的图像
        if not skip_test:
            render_set(dataset.model_path, "test", scene.loaded_iter, scene.getTestCameras(), gaussians, pipeline, background)

但是这两个方法都是外层函数，并没有展示渲染如何进行参数传递和具体操作，在以下代码中才是最关键的内容：

gaussian_renderer\__init__.py
line 18
def render(viewpoint_camera, pc: GaussianModel, pipe, bg_color: torch.Tensor, scaling_modifier=1.0, override_color=None):
    """
    渲染场景。
    
    参数:
    viewpoint_camera - 摄像机视角
    pc - 高斯模型
    pipe - 管道参数
    bg_color - 背景颜色张量，必须在GPU上
    scaling_modifier - 缩放修饰符，默认为1.0
    override_color - 覆盖颜色，默认为None
    """
 
    # 创建一个全零张量，用于使PyTorch返回2D（屏幕空间）均值的梯度
    screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
    try:
        screenspace_points.retain_grad()  # 保留梯度信息
    except:
        pass

    # 设置光栅化配置
    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)  # 计算视角的X轴正切
    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)  # 计算视角的Y轴正切

    raster_settings = GaussianRasterizationSettings(
        image_height=int(viewpoint_camera.image_height),  # 图像高度
        image_width=int(viewpoint_camera.image_width),  # 图像宽度
        tanfovx=tanfovx,  # 视角X轴正切
        tanfovy=tanfovy,  # 视角Y轴正切
        bg=bg_color,  # 背景颜色
        scale_modifier=scaling_modifier,  # 缩放修饰符
        viewmatrix=viewpoint_camera.world_view_transform,  # 世界视图变换矩阵
        projmatrix=viewpoint_camera.full_proj_transform,  # 投影变换矩阵
        sh_degree=pc.active_sh_degree,  # 球谐函数度数
        campos=viewpoint_camera.camera_center,  # 摄像机中心
        prefiltered=False,  # 预过滤
        debug=pipe.debug  # 调试模式
    )

    rasterizer = GaussianRasterizer(raster_settings=raster_settings)  # 初始化光栅化器

    means3D = pc.get_xyz  # 获取3D均值
    means2D = screenspace_points  # 获取2D均值
    opacity = pc.get_opacity  # 获取不透明度

    # 如果提供了预计算的3D协方差，则使用它。如果没有，则从光栅化器的缩放/旋转中计算。
    scales = None
    rotations = None
    cov3D_precomp = None
    if pipe.compute_cov3D_python:
        cov3D_precomp = pc.get_covariance(scaling_modifier)  # 计算3D协方差
    else:
        scales = pc.get_scaling  # 获取缩放
        rotations = pc.get_rotation  # 获取旋转

    # 如果提供了预计算的颜色，则使用它们。否则，如果需要在Python中预计算SH到颜色的转换，则进行转换。
    shs = None
    colors_precomp = None
    if override_color is None:
        if pipe.convert_SHs_python:
            shs_view = pc.get_features.transpose(1, 2).view(-1, 3, (pc.max_sh_degree + 1) ** 2)
            dir_pp = (pc.get_xyz - viewpoint_camera.camera_center.repeat(pc.get_features.shape[0], 1))
            dir_pp_normalized = dir_pp / dir_pp.norm(dim=1, keepdim=True)
            sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized)
            colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)  # 计算颜色
        else:
            shs = pc.get_features  # 获取球谐函数特征
    else:
        colors_precomp = override_color  # 覆盖颜色

    # 将可见的高斯体光栅化为图像，并获取它们在屏幕上的半径。
    rendered_image, radii = rasterizer(
        means3D=means3D,
        means2D=means2D,
        shs=shs,
        colors_precomp=colors_precomp,
        opacities=opacity,
        scales=scales,
        rotations=rotations,
        cov3D_precomp=cov3D_precomp)

    # 那些被视锥剔除或半径为0的高斯体是不可见的。
    # 它们将被排除在用于分裂标准的值更新之外。
    return {
        "render": rendered_image,  # 渲染图像
        "viewspace_points": screenspace_points,  # 视图空间点
        "visibility_filter": radii > 0,  # 可见性过滤器
        "radii": radii  # 半径
    }

最值得关注的光栅化器，如果转到定义去查看，其实会发现它就是第二期里讲forward的代码，只是这里面用python写了变量的调用，实际的操作方式还是在cu文件里面。所以在此就不多做赘述，可以看上一期博客里面对forward的解读。