3D Gaussian Splatting代码详解（三）：模型构建，实现3D 高斯椭球体的克隆和分裂

3 模型构建

3.4 根据梯度对3D gaussian 进行增加或删减

（1）对3D高斯分布进行密集化和修剪的操作

def densify_and_prune(self, max_grad, min_opacity, extent, max_screen_size):"""对3D高斯分布进行密集化和修剪的操作:param max_grad: 梯度的最大阈值，用于判断是否需要克隆或分割。:param min_opacity: 不透明度的最小阈值，低于此值的3D高斯将被删除。:param extent: 场景的尺寸范围，用于评估高斯分布的大小是否合适。:param max_screen_size: 最大屏幕尺寸阈值，用于修剪过大的高斯分布。"""# 计算3D高斯中心的累积梯度并修正NaN值grads = self.xyz_gradient_accum / self.denomgrads[grads.isnan()] = 0.0# 根据梯度和尺寸阈值进行克隆或分割操作self.densify_and_clone(grads, max_grad, extent)self.densify_and_split(grads, max_grad, extent)# 创建修剪掩码以删除不必要的3D高斯分布prune_mask = (self.get_opacity &lt; min_opacity).squeeze()if max_screen_size:big_points_vs = self.max_radii2D &gt; max_screen_sizebig_points_ws = self.get_scaling.max(dim=1).values &gt; 0.1 * extentprune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws)# 应用修剪掩码self.prune_points(prune_mask)# 清理CUDA缓存以释放资源torch.cuda.empty_cache()

这是一个用于密集化和修剪3D高斯分布的函数 densify_and_prune，其作用是根据梯度、不透明度、尺寸范围以及屏幕尺寸等参数来动态调整和优化3D高斯分布。

函数参数

max_grad: 定义梯度的最大阈值，用于判断高斯分布是否需要进一步克隆或分割。
min_opacity: 定义不透明度的最小阈值，低于该值的高斯分布将被删除。
extent: 定义场景的总体尺寸范围，用来评估高斯分布的大小是否在合理范围内。
max_screen_size: 定义屏幕上的最大显示尺寸阈值，用于修剪尺寸过大的高斯分布。

功能描述

计算梯度并处理 NaN 值：

函数首先计算每个3D高斯中心点的累积梯度，并将任何 NaN 值替换为 0，以避免异常值影响密集化过程。

密集化操作：

使用 densify_and_clone 和 densify_and_split 函数，根据梯度和尺寸范围的条件，来进行高斯分布的克隆和分割操作，使得数据表示更加密集，并有利于提高信息捕获效率。

修剪掩码的创建：

生成一个 prune_mask 掩码，根据不透明度阈值 min_opacity 来筛选出不符合条件的3D高斯。
若 max_screen_size 存在，函数会额外考虑屏幕尺寸阈值，确保尺寸过大的高斯点也会被修剪掉。这里 big_points_vs 和 big_points_ws 用于检测高斯在屏幕和世界坐标下的尺寸是否超过限制。

应用修剪掩码：

使用 prune_points 方法删除符合掩码条件的3D高斯分布，减少冗余并优化性能。

清理 CUDA 缓存：

最后，通过 torch.cuda.empty_cache() 清理 GPU 缓存，释放未使用的显存资源，以确保内存占用最小化并提高计算效率。

此函数适用于3D场景的稠密表示和资源优化，通过密集化和修剪操作，有效提升了3D高斯分布的适应性和计算性能，适合动态场景中的高效数据处理。

（2）删除不符合要求的3D高斯分布

def prune_points(self, mask):"""删除不符合要求的3D高斯分布。:param mask: 一个布尔张量，表示需要删除的3D高斯分布。"""# 生成有效点的掩码并更新优化器中的参数valid_points_mask = ~maskoptimizable_tensors = self._prune_optimizer(valid_points_mask)# 更新各参数self._xyz = optimizable_tensors["xyz"]self._features_dc = optimizable_tensors["f_dc"]self._features_rest = optimizable_tensors["f_rest"]self._opacity = optimizable_tensors["opacity"]self._scaling = optimizable_tensors["scaling"]self._rotation = optimizable_tensors["rotation"]# 更新累积梯度和其他相关张量self.xyz_gradient_accum = self.xyz_gradient_accum[valid_points_mask]self.denom = self.denom[valid_points_mask]self.max_radii2D = self.max_radii2D[valid_points_mask]

（3）删除不符合要求的3D高斯分布在优化器中对应的参数

def _prune_optimizer(self, mask):"""删除不符合要求的3D高斯分布在优化器中对应的参数:param mask: 一个布尔张量，表示需要保留的3D高斯分布。:return: 更新后的可优化张量字典。"""optimizable_tensors = {}for group in self.optimizer.param_groups:stored_state = self.optimizer.state.get(group['params'][0], None)if stored_state is not None:# 更新优化器状态stored_state["exp_avg"] = stored_state["exp_avg"][mask]stored_state["exp_avg_sq"] = stored_state["exp_avg_sq"][mask]# 删除旧状态并更新参数del self.optimizer.state[group['params'][0]]group["params"][0] = nn.Parameter((group["params"][0][mask].requires_grad_(True)))self.optimizer.state[group['params'][0]] = stored_stateoptimizable_tensors[group["name"]] = group["params"][0]else:group["params"][0] = nn.Parameter(group["params"][0][mask].requires_grad_(True))optimizable_tensors[group["name"]] = group["params"][0]return optimizable_tensors

（4）对那些梯度超过一定阈值且尺度小于一定阈值的3D高斯进行克隆操作。

def densify_and_clone(self, grads, grad_threshold, scene_extent):"""对那些梯度超过一定阈值且尺度小于一定阈值的3D高斯进行克隆操作。这意味着这些高斯在空间中可能表示的细节不足，需要通过克隆来增加细节。"""# 选择满足条件的点selected_pts_mask = torch.where(torch.norm(grads, dim=-1) &gt;= grad_threshold, True, False)selected_pts_mask = torch.logical_and(selected_pts_mask,torch.max(self.get_scaling, dim=1).values &lt;= self.percent_dense * scene_extent)# 提取这些点的属性new_xyz = self._xyz[selected_pts_mask]  # 位置new_features_dc = self._features_dc[selected_pts_mask]  # 直流分量（基本颜色）new_features_rest = self._features_rest[selected_pts_mask]  # 其他球谐分量new_opacities = self._opacity[selected_pts_mask]  # 不透明度new_scaling = self._scaling[selected_pts_mask]  # 尺度new_rotation = self._rotation[selected_pts_mask]  # 旋转# 将克隆得到的新高斯分布的属性添加到模型中self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation)

（5）对那些梯度超过一定阈值且尺度大于一定阈值的3D高斯进行分割操作。

def densify_and_split(self, grads, grad_threshold, scene_extent, N=2):"""对那些梯度超过一定阈值且尺度大于一定阈值的3D高斯进行分割操作。这意味着这些高斯可能过于庞大，覆盖了过多的空间区域，需要分割成更小的部分以提升细节。"""# 初始化n_init_points = self.get_xyz.shape[0]padded_grad = torch.zeros((n_init_points), device="cuda")padded_grad[:grads.shape[0]] = grads.squeeze()# 选择满足条件的点selected_pts_mask = torch.where(padded_grad &gt;= grad_threshold, True, False)selected_pts_mask = torch.logical_and(selected_pts_mask,torch.max(self.get_scaling, dim=1).values &gt; self.percent_dense * scene_extent)# 计算新高斯分布的属性stds = self.get_scaling[selected_pts_mask].repeat(N, 1)  # 尺度means = torch.zeros((stds.size(0), 3), device="cuda")  # 均值（新分布的中心点）samples = torch.normal(mean=means, std=stds)  # 随机采样新的位置rots = build_rotation(self._rotation[selected_pts_mask]).repeat(N, 1, 1)  # 旋转# 计算新的位置new_xyz = torch.bmm(rots, samples.unsqueeze(-1)).squeeze(-1) + self.get_xyz[selected_pts_mask].repeat(N, 1)# 调整尺度并保持其他属性new_scaling = self.scaling_inverse_activation(self.get_scaling[selected_pts_mask].repeat(N, 1) / (0.8 * N))new_rotation = self._rotation[selected_pts_mask].repeat(N, 1)new_features_dc = self._features_dc[selected_pts_mask].repeat(N, 1, 1)new_features_rest = self._features_rest[selected_pts_mask].repeat(N, 1, 1)new_opacity = self._opacity[selected_pts_mask].repeat(N, 1)# 将分割得到的新高斯分布的属性添加到模型中self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacity, new_scaling, new_rotation)# 删除原有过大的高斯分布prune_filter = torch.cat((selected_pts_mask, torch.zeros(N * selected_pts_mask.sum(), device="cuda", dtype=bool)))self.prune_points(prune_filter)

（5）将新生成的3D高斯分布的属性添加到模型的参数中。

def densification_postfix(self, new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation):"""将新生成的3D高斯分布的属性添加到模型的参数中。"""d = {"xyz": new_xyz,"f_dc": new_features_dc,"f_rest": new_features_rest,"opacity": new_opacities,"scaling": new_scaling,"rotation": new_rotation}optimizable_tensors = self.cat_tensors_to_optimizer(d)self._xyz = optimizable_tensors["xyz"]self._features_dc = optimizable_tensors["f_dc"]self._features_rest = optimizable_tensors["f_rest"]self._opacity = optimizable_tensors["opacity"]self._scaling = optimizable_tensors["scaling"]self._rotation = optimizable_tensors["rotation"]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")self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")

（6）将新的参数张量添加到优化器的参数组中

def cat_tensors_to_optimizer(self, tensors_dict):"""将新的参数张量添加到优化器的参数组中"""optimizable_tensors = {}for group in self.optimizer.param_groups:extension_tensor = tensors_dict[group["name"]]stored_state = self.optimizer.state.get(group['params'][0], None)if stored_state is not None:stored_state["exp_avg"] = torch.cat((stored_state["exp_avg"], torch.zeros_like(extension_tensor)), dim=0)stored_state["exp_avg_sq"] = torch.cat((stored_state["exp_avg_sq"], torch.zeros_like(extension_tensor)), dim=0)del self.optimizer.state[group['params'][0]]group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))self.optimizer.state[group['params'][0]] = stored_stateoptimizable_tensors[group["name"]] = group["params"][0]else:group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))optimizable_tensors[group["name"]] = group["params"][0]return optimizable_tensors

（7）根据旋转四元数构建旋转矩阵。

def build_rotation(r):"""根据旋转四元数构建旋转矩阵。"""norm = torch.sqrt(r[:,0]*r[:,0] + r[:,1]*r[:,1] + r[:,2]*r[:,2] + r[:,3]*r[:,3])q = r / norm[:, None]R = torch.zeros((q.size(0), 3, 3), device='cuda')# 计算旋转矩阵的各元素...return R