3D Gaussian Splatting代码详解（二）：模型构建

3 模型构建

gaussians = GaussianModel(dataset.sh_degree)

3.1 初始化函数

__init__ 构造函数

构造函数 __init__ 的主要作用是初始化 3D 高斯模型的各项参数和激活函数，用于生成 3D 空间中的高斯表示。

1. 初始化球谐函数的参数：

   • self.active_sh_degree：当前激活的球谐函数次数，初始为 0。
   • self.max_sh_degree：指定的最大球谐函数次数（由 sh_degree 参数传入），用于控制颜色表示的复杂度。

2. 初始化 3D 高斯模型的各项参数：

   • self._xyz：3D 高斯中心位置（均值），用于确定高斯分布的中心。
   • self._features_dc：第一个球谐系数，用于表示基础颜色。
   • self._features_rest：其余球谐系数，用于表示颜色细节。
   • self._scaling：3D 高斯的尺度参数，控制高斯分布的宽度。
   • self._rotation：3D 高斯的旋转参数，使用四元数表示旋转。
   • self._opacity：3D 高斯的不透明度，控制可见性。
   • self.max_radii2D：在 2D 投影中的最大半径，控制每个高斯点在屏幕上呈现的大小。
   • self.xyz_gradient_accum：用于累积 3D 高斯中心位置的梯度信息，可能在优化时会用到。
   • self.denom：未明确用途的参数，但可能是为后续的计算准备。
   • self.optimizer：优化器，将用于调整上述参数以优化模型效果。

3. 调用 setup_functions 初始化处理函数 ：setup_functions 定义了一些用于操作这些参数的处理函数（如协方差矩阵构建、激活函数等）。

setup_functions 函数

setup_functions 是一个辅助方法，用于定义和初始化一些与 3D 高斯模型参数有关的函数。这些函数主要用于 3D 高斯模型的计算和优化，特别是协方差矩阵、激活函数等。

1. 构建协方差矩阵：

   • build_covariance_from_scaling_rotation：根据高斯的尺度参数和旋转参数构建协方差矩阵。
     • 使用 build_scaling_rotation（假设是外部函数）对 scaling_modifier * scaling 和 rotation 参数进行处理，生成一个变换矩阵 L。
     • 计算实际协方差矩阵 actual_covariance = L @ L.transpose(1, 2)，从而得到对称的协方差矩阵 symm。

2. 激活函数：

   • self.scaling_activation 和 self.scaling_inverse_activation：分别为尺度参数的激活函数（torch.exp）和逆激活函数（torch.log），确保尺度参数为非负。
   • self.covariance_activation：用于计算协方差矩阵的激活函数，设置为 build_covariance_from_scaling_rotation。
   • self.opacity_activation 和 self.inverse_opacity_activation：分别为不透明度的激活函数（torch.sigmoid）和逆激活函数（inverse_sigmoid），确保不透明度在 0 到 1 之间。
   • self.rotation_activation：用于标准化旋转参数的函数，使用 torch.nn.functional.normalize 确保四元数的单位长度，维持旋转参数的有效性。

代码功能：这段代码为 3D 高斯模型定义了各种属性和函数，能够支持对模型位置、颜色、透明度、旋转等特征进行控制，并提供协方差矩阵生成、尺度和不透明度控制等方法。

代码如下：

def __init__(self, sh_degree: int):
    """
    初始化3D高斯模型的参数。

    :param sh_degree: 球谐函数的最大次数，用于控制颜色表示的复杂度。
    """
    # 初始化球谐次数和最大球谐次数
    self.active_sh_degree = 0  # 当前激活的球谐次数，初始为0
    self.max_sh_degree = sh_degree  # 允许的最大球谐次数

    # 初始化3D高斯模型的各项参数
    self._xyz = torch.empty(0)  # 3D高斯的中心位置（均值）
    self._features_dc = torch.empty(0)  # 第一个球谐系数，用于表示基础颜色
    self._features_rest = torch.empty(0)  # 其余的球谐系数，用于表示颜色的细节和变化
    self._scaling = torch.empty(0)  # 3D高斯的尺度参数，控制高斯的宽度
    self._rotation = torch.empty(0)  # 3D高斯的旋转参数，用四元数表示
    self._opacity = torch.empty(0)  # 3D高斯的不透明度，控制可见性
    self.max_radii2D = torch.empty(0)  # 在2D投影中，每个高斯的最大半径
    self.xyz_gradient_accum = torch.empty(0)  # 用于累积3D高斯中心位置的梯度
    self.denom = torch.empty(0)  # 未明确用途的参数
    self.optimizer = None  # 优化器，用于调整上述参数以改进模型

    # 调用setup_functions来初始化一些处理函数
    self.setup_functions()

    def setup_functions(self):
    """
    定义和初始化一些用于处理3D高斯模型参数的函数。
    """
    
    # 定义构建3D高斯协方差矩阵的函数
    def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation):
        L = build_scaling_rotation(scaling_modifier * scaling, rotation)
        actual_covariance = L @ L.transpose(1, 2)  # 计算实际的协方差矩阵
        symm = strip_symmetric(actual_covariance)  # 提取对称部分
        return symm

    # 初始化一些激活函数
    self.scaling_activation = torch.exp  # 用exp函数确保尺度参数非负
    self.scaling_inverse_activation = torch.log  # 尺度参数的逆激活函数，用于梯度回传

    self.covariance_activation = build_covariance_from_scaling_rotation  # 协方差矩阵的激活函数

    self.opacity_activation = torch.sigmoid  # 用sigmoid函数确保不透明度在0到1之间
    self.inverse_opacity_activation = inverse_sigmoid  # 不透明度的逆激活函数

    self.rotation_activation = torch.nn.functional.normalize  # 用于标准化旋转参数的函数
	

---

def build_scaling_rotation(s, r):
    """
    构建3D高斯模型的尺度-旋转矩阵。

    :param s: 尺度参数。
    :param r: 旋转参数。
    :return: 尺度-旋转矩阵。
    """
    L = torch.zeros((s.shape[0], 3, 3), dtype=torch.float, device="cuda")  # 初始化尺度矩阵
    R = build_rotation(r)  # 计算旋转矩阵

    # 设置尺度矩阵的对角线元素
    L[:, 0, 0] = s[:, 0]
    L[:, 1, 1] = s[:, 1]
    L[:, 2, 2] = s[:, 2]

    L = R @ L  # 应用旋转
    return L

def strip_symmetric(sym):
    """
    提取协方差矩阵的对称部分。

    :param sym: 协方差矩阵。
    :return: 对称部分。
    """
    return strip_lowerdiag(sym)

def strip_lowerdiag(L):
    """
    从协方差矩阵中提取六个独立参数。

    :param L: 协方差矩阵。
    :return: 六个独立参数组成的张量。
    """
    uncertainty = torch.zeros((L.shape[0], 6), dtype=torch.float, device="cuda")

    # 提取协方差矩阵的独立元素
    uncertainty[:, 0] = L[:, 0, 0]
    uncertainty[:, 1] = L[:, 0, 1]
    uncertainty[:, 2] = L[:, 0, 2]
    uncertainty[:, 3] = L[:, 1, 1]
    uncertainty[:, 4] = L[:, 1, 2]
    uncertainty[:, 5] = L[:, 2, 2]
    return uncertainty

分别用于构建 3D 高斯模型的尺度-旋转矩阵，并从协方差矩阵中提取其对称部分和独立参数。这一系列函数主要用于 3D 空间中的高斯分布建模，特别是通过旋转和缩放来控制高斯的形状和方向。

函数解析

1. build_scaling_rotation(s, r)

build_scaling_rotation 的作用是生成 3D 高斯模型的尺度-旋转矩阵，方法是首先创建一个对角尺度矩阵 LLL ，然后对其应用旋转矩阵 RRR。

• 参数：

  • s: 尺度参数，形状为 (batch_size, 3) 的张量，表示每个样本在 3 个维度上的尺度（宽度）。
  • r: 旋转参数，通常以四元数或欧拉角表示，用于生成 旋转矩阵 RRR。

• 步骤：

  1. 初始化张量 L：一个形状为 (batch_size, 3, 3) 的 3D 张量，全为零。
  2. 计算旋转矩阵 R：调用 build_rotation(r) 函数（假设在代码的其他部分定义），生成一个形状为 (batch_size, 3, 3) 的旋转矩阵。
  3. 设置尺度矩阵的对角元素：将 s 的每个维度赋值到 L 的对角线上（分别为 L[:, 0, 0], L[:, 1, 1], L[:, 2, 2]）。
  4. 计算尺度-旋转矩阵：将旋转矩阵 R 与尺度矩阵 L 相乘，以应用旋转变换，得到最终的尺度-旋转矩阵 L。

• 返回值 ：返回 L，即应用旋转后的尺度矩阵，形状为 (batch_size, 3, 3)。

2. strip_symmetric(sym)

strip_symmetric 函数用于提取协方差矩阵的对称部分。协方差矩阵通常是对称的，这一函数可以用来获取对称部分。

• 参数：

  • sym: 协方差矩阵，形状为 (batch_size, 3, 3)。

• 步骤：

  1. 调用 strip_lowerdiag(sym)，该函数提取协方差矩阵中的 6 个独立元素（下三角部分）来表示对称矩阵。

• 返回值：返回协方差矩阵的对称部分的六个独立参数。

3. strip_lowerdiag(L)

strip_lowerdiag 用于从协方差矩阵 L 中提取独立参数（即下三角部分的 6 个独立值），表示协方差矩阵的对称部分。这样可以用更少的参数来表示矩阵。

• 参数：

  • L: 协方差矩阵，形状为 (batch_size, 3, 3)。

• 步骤：

  1. 初始化 uncertainty：形状为 (batch_size, 6) 的张量，用来存储每个协方差矩阵的独立参数。
  2. 提取下三角元素：分别提取对角和非对角元素，存入 uncertainty 的不同位置。
     • uncertainty[:, 0]：取 L 的第一个对角元素 L[:, 0, 0]。
     • uncertainty[:, 1] 和 uncertainty[:, 2]：分别取 L[:, 0, 1] 和 L[:, 0, 2]。
     • uncertainty[:, 3]：取 L 的第二个对角元素 L[:, 1, 1]。
     • uncertainty[:, 4]：取 L[:, 1, 2]。
     • uncertainty[:, 5]：取 L 的第三个对角元素 L[:, 2, 2]。

• 返回值 ：返回 uncertainty 张量，包含六个独立参数，形状为 (batch_size, 6)。

代码总结

这段代码构建了 3D 高斯模型的尺度-旋转矩阵，并通过提取协方差矩阵的对称部分，简化了矩阵的存储和计算。 build_scaling_rotation 通过结合旋转和缩放矩阵来生成高斯的尺度-旋转矩阵，而 strip_lowerdiag 和 strip_symmetric 则用于提取协方差矩阵的主要参数。这种方法在 3D 高斯建模中常用于表征形状、方向和尺度，减少了参数数量，提高了效率。

3.2 为3D Gaussion的各组参数创建 optimizer 以及lr_scheduler

def training_setup(self, training_args):
    """
    设置训练参数，包括初始化用于累积梯度的变量，配置优化器，以及创建学习率调度器

    :param training_args: 包含训练相关参数的对象。
    """
    # 设置在训练过程中，用于密集化处理的3D高斯点的比例
    self.percent_dense = training_args.percent_dense

    # 初始化用于累积3D高斯中心点位置梯度的张量，用于之后判断是否需要对3D高斯进行克隆或切分
    self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")

    self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")

    # 配置各参数的优化器，包括指定参数、学习率和参数名称
    l = [
        {'params': [self._xyz], 'lr': training_args.position_lr_init * self.spatial_lr_scale, "name": "xyz"},
        {'params': [self._features_dc], 'lr': training_args.feature_lr, "name": "f_dc"},
        {'params': [self._features_rest], 'lr': training_args.feature_lr / 20.0, "name": "f_rest"},
        {'params': [self._opacity], 'lr': training_args.opacity_lr, "name": "opacity"},
        {'params': [self._scaling], 'lr': training_args.scaling_lr, "name": "scaling"},
        {'params': [self._rotation], 'lr': training_args.rotation_lr, "name": "rotation"}
    ]
    # 创建优化器，这里使用Adam优化器
    self.optimizer = torch.optim.Adam(l, lr=0.0, eps=1e-15)

    # 创建学习率调度器，用于对中心点位置的学习率进行调整
    self.xyz_scheduler_args = get_expon_lr_func(
        lr_init=training_args.position_lr_init*self.spatial_lr_scale,
        lr_final=training_args.position_lr_final*self.spatial_lr_scale,
        lr_delay_mult=training_args.position_lr_delay_mult,
        max_steps=training_args.position_lr_max_steps
    )

def get_expon_lr_func(lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000):
    """
    创建一个学习率调度函数，该函数根据训练进度动态调整学习率

    :param lr_init: 初始学习率。
    :param lr_final: 最终学习率。
    :param lr_delay_steps: 学习率延迟步数，在这些步数内学习率将被降低。
    :param lr_delay_mult: 学习率延迟乘数，用于计算初始延迟学习率。
    :param max_steps: 最大步数，用于规范化训练进度。
    :return: 一个函数，根据当前步数返回调整后的学习率。
    """
    def helper(step):
        # 如果步数小于0或学习率为0，直接返回0，表示不进行优化
        if step < 0 or (lr_init == 0.0 and lr_final == 0.0):
            return 0.0
        # 如果设置了学习率延迟步数，计算延迟调整后的学习率
        if lr_delay_steps > 0:
            delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin(
                0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1)
            )
        else:
            delay_rate = 1.0
        # 根据步数计算学习率的对数线性插值，实现从初始学习率到最终学习率的平滑过渡
        t = np.clip(step / max_steps, 0, 1)
        log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t)
        # 返回调整后的学习率
        return delay_rate * log_lerp

    return helper

3.3 从点云pcd 创建3D Gaussian

def create_from_pcd(self, pcd: BasicPointCloud, spatial_lr_scale: float):
    """
    从点云数据初始化模型参数。

    :param pcd: 点云数据，包含点的位置和颜色。
    :param spatial_lr_scale: 空间学习率缩放因子，影响位置参数的学习率。
    """
    # 将点云的位置和颜色数据从numpy数组转换为PyTorch张量，并传送到CUDA设备上
    self.spatial_lr_scale = spatial_lr_scale
    fused_point_cloud = torch.tensor(np.asarray(pcd.points)).float().cuda()  # (P, 3)
    fused_color = RGB2SH(torch.tensor(np.asarray(pcd.colors)).float().cuda())  # (P, 3)

    # 初始化存储球谐系数的张量，每个颜色通道有(max_sh_degree + 1) ** 2个球谐系数
    features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda()  # (P, 3, 16)
    features[:, :3, 0] = fused_color  # 将RGB转换后的球谐系数C0项的系数存入
    features[:, 3:, 1:] = 0.0  # 其余球谐系数初始化为0

    # 打印初始点的数量
    print("Number of points at initialisation : ", fused_point_cloud.shape[0])

    # 计算点云中每个点到其最近的k个点的平均距离的平方，用于确定高斯的尺度参数
    dist2 = torch.clamp_min(distCUDA2(torch.from_numpy(np.asarray(pcd.points)).float().cuda()), 0.0000001)  # (P,)
    scales = torch.log(torch.sqrt(dist2))[..., None].repeat(1, 3)  # (P, 3)

    # 初始化每个点的旋转参数为单位四元数（无旋转）
    rots = torch.zeros((fused_point_cloud.shape[0], 4), device="cuda")  # (P, 4)
    rots[:, 0] = 1  # 四元数的实部为1，表示无旋转

    # 初始化每个点的不透明度为0.1（通过inverse_sigmoid转换）
    opacities = inverse_sigmoid(0.1 * torch.ones((fused_point_cloud.shape[0], 1), dtype=torch.float, device="cuda"))  # (P, 1)

    # 将以上计算的参数设置为模型的可训练参数
    self._xyz = nn.Parameter(fused_point_cloud.requires_grad_(True))  # 位置
    self._features_dc = nn.Parameter(features[:, :, 0:1].transpose(1, 2).contiguous().requires_grad_(True))  # 球谐系数C0项
    self._features_rest = nn.Parameter(features[:, :, 1:].transpose(1, 2).contiguous().requires_grad_(True))  # 其余球谐系数
    self._scaling = nn.Parameter(scales.requires_grad_(True))  # 尺度
    self._rotation = nn.Parameter(rots.requires_grad_(True))  # 旋转
    self._opacity = nn.Parameter(opacities.requires_grad_(True))  # 不透明度
    self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")  # 存储2D投影的最大半径，初始化为0

def RGB2SH(rgb):
    """
    将RGB颜色值转换为球谐系数C0项的系数。

    :param rgb: RGB颜色值。
    :return: 转换后的球谐系数C0项的系数。
    """
    return (rgb - 0.5) / C0