3D Gaussian Splatting代码详解(二):模型构建
3 模型构建
gaussians = GaussianModel(dataset.sh_degree)
3.1 初始化函数
__init__
构造函数
构造函数 __init__
的主要作用是初始化 3D 高斯模型的各项参数和激活函数,用于生成 3D 空间中的高斯表示。
初始化球谐函数的参数:
self.active_sh_degree
:当前激活的球谐函数次数,初始为0
。self.max_sh_degree
:指定的最大球谐函数次数(由sh_degree
参数传入),用于控制颜色表示的复杂度。初始化 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
:优化器,将用于调整上述参数以优化模型效果。调用
setup_functions
初始化处理函数:setup_functions
定义了一些用于操作这些参数的处理函数(如协方差矩阵构建、激活函数等)。
setup_functions
函数
setup_functions
是一个辅助方法,用于定义和初始化一些与 3D 高斯模型参数有关的函数。这些函数主要用于 3D 高斯模型的计算和优化,特别是协方差矩阵、激活函数等。
构建协方差矩阵:
build_covariance_from_scaling_rotation
:根据高斯的尺度参数和旋转参数构建协方差矩阵。
- 使用
build_scaling_rotation
(假设是外部函数)对scaling_modifier * scaling
和rotation
参数进行处理,生成一个变换矩阵L
。- 计算实际协方差矩阵
actual_covariance = L @ L.transpose(1, 2)
,从而得到对称的协方差矩阵symm
。激活函数:
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。
-
步骤:
- 初始化张量
L
:一个形状为(batch_size, 3, 3)
的 3D 张量,全为零。 - 计算旋转矩阵
R
:调用build_rotation(r)
函数(假设在代码的其他部分定义),生成一个形状为(batch_size, 3, 3)
的旋转矩阵。 - 设置尺度矩阵的对角元素:将
s
的每个维度赋值到L
的对角线上(分别为L[:, 0, 0]
,L[:, 1, 1]
,L[:, 2, 2]
)。 - 计算尺度-旋转矩阵:将旋转矩阵
R
与尺度矩阵L
相乘,以应用旋转变换,得到最终的尺度-旋转矩阵L
。
- 初始化张量
-
返回值:返回
L
,即应用旋转后的尺度矩阵,形状为(batch_size, 3, 3)
。
2. strip_symmetric(sym)
strip_symmetric
函数用于提取协方差矩阵的对称部分。协方差矩阵通常是对称的,这一函数可以用来获取对称部分。
-
参数:
sym
: 协方差矩阵,形状为(batch_size, 3, 3)
。
-
步骤:
- 调用
strip_lowerdiag(sym)
,该函数提取协方差矩阵中的 6 个独立元素(下三角部分)来表示对称矩阵。
- 调用
-
返回值:返回协方差矩阵的对称部分的六个独立参数。
3. strip_lowerdiag(L)
strip_lowerdiag
用于从协方差矩阵 L
中提取独立参数(即下三角部分的 6 个独立值),表示协方差矩阵的对称部分。这样可以用更少的参数来表示矩阵。
-
参数:
L
: 协方差矩阵,形状为(batch_size, 3, 3)
。
-
步骤:
- 初始化
uncertainty
:形状为(batch_size, 6)
的张量,用来存储每个协方差矩阵的独立参数。 - 提取下三角元素:分别提取对角和非对角元素,存入
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