3D Gaussian Splatting代码详解(二):模型构建

发布于:2024-11-02 ⋅ 阅读:(13) ⋅ 点赞:(0)

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 * scalingrotation 参数进行处理,生成一个变换矩阵 L
      • 计算实际协方差矩阵 actual_covariance = L @ L.transpose(1, 2),从而得到对称的协方差矩阵 symm
  2. 激活函数

    • self.scaling_activationself.scaling_inverse_activation:分别为尺度参数的激活函数(torch.exp)和逆激活函数(torch.log),确保尺度参数为非负。
    • self.covariance_activation:用于计算协方差矩阵的激活函数,设置为 build_covariance_from_scaling_rotation
    • self.opacity_activationself.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_lowerdiagstrip_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