Splatstudio 学习笔记

NeRF_er

已于 2024-09-06 19:15:17 修改

阅读量880

点赞数 8

文章标签：学习笔记

于 2024-04-08 16:25:15 首次发布

本文链接：https://blog.csdn.net/qq_41623632/article/details/137477376

版权

1. 3DGS 是全图的Render, 因此特地写了 full_images_datamanger.py 每个step 取出一张图像。

返回值是一张全图的 RGB 和对应的 Camera

2. 3D GS 没有生成光线，因此不需要指定near 和 far，即 collider是None。但需要对3D 高斯球进行初始化：

高斯球的means (中心点的初始化)：输入点云的坐标当作是高斯球的中心点：

 means = torch.nn.Parameter(self.seed_points[0])  # (Location, Color)

高斯球的scale初始化，使用KNN 算法计算每一个点云到最近3个点之间的距离 distance, 对这个距离求解平均距离，以这个 ave_distance 的 log 对数的结果当作是 scale 的初始化：

distances, _ = self.k_nearest_sklearn(means.data, 3)
distances = torch.from_numpy(distances)
# find the average of the three nearest neighbors for each point and use that as the scale
avg_dist = distances.mean(dim=-1, keepdim=True)
scales = torch.nn.Parameter(torch.log(avg_dist.repeat(1, 3)))

旋转四元数的初始化，随机初始化：

quats = torch.nn.Parameter(random_quat_tensor(num_points))

opacity的初始化，随机给定一个固定数值从 0.1 初始化：

opacities = torch.nn.Parameter(torch.logit(0.1 * torch.ones(num_points, 1)))

color 的SH 参数的初始化，从点云的 RGB 直接转化为 SH 系数：

if self.config.sh_degree > 0:
   shs[:, 0, :3] = RGB2SH(self.seed_points[1] / 255) ## RGB 转化成第0阶的 SH 系数
   shs[:, 1:, 3:] = 0.0
else:
   CONSOLE.log("use color only optimization with sigmoid activation")
   shs[:, 0, :3] = torch.logit(self.seed_points[1] / 255, eps=1e-10)
features_dc = torch.nn.Parameter(shs[:, 0, :])
features_rest = torch.nn.Parameter(shs[:, 1:, :])

SH 的稀疏的 shape 标准是（N，16，3）； 3 代表3个通道， N 代表 N个点，16 代表阶数。

第0 阶的参数直接由RGB 转换而来；之后都是设置为0 进行初始化。对应的优化参数是：

features_dc = torch.nn.Parameter(shs[:, 0, :]) ## 【N,3】
features_rest = torch.nn.Parameter(shs[:, 1:, :]) ## 【N,15,3】

3. 3D GS 的属性：

3D GS 如果没有输入点云，那么使用的是 torc.random 函数进行随机初始化。

 self.gauss_params = torch.nn.ParameterDict(
            {
                "means": means,  ## Location
                "scales": scales, ## 缩放因子
                "quats": quats,   ## 旋转的四元数
                "features_dc": features_dc,
                "features_rest": features_rest,
                "opacities": opacities, ## 不透明度
            }
        )

这行使用 torch 的 ParameterDict 的代码相当于直接使用：

self.means = means;
self.scales = scales;

3.1 首先对于整张图像进行 DownScale 4 倍：

对图像进行4倍降采样,并修改对应的内参
camera_downscale = self._get_downscale_factor() ## 4
camera.rescale_output_resolution(1 / camera_downscale)

3.2 修改 Camera 的坐标，将其从 nerfstudio 的坐标旋转到 opencv 的坐标。也就是 [-1,-1,1] 的矩阵

 # shift the camera to center of scene looking at center
R = camera.camera_to_worlds[0, :3, :3]  # 3 x 3
T = camera.camera_to_worlds[0, :3, 3:4]  # 3 x 1
# flip the z and y axes to align with gsplat conventions
R_edit = torch.diag(torch.tensor([1, -1, -1], device=self.device, dtype=R.dtype))
R = R @ R_edit
# analytic matrix inverse to get world2camera matrix
R_inv = R.T
T_inv = -R_inv @ T
viewmat = torch.eye(4, device=R.device, dtype=R.dtype)
viewmat[:3, :3] = R_inv
viewmat[:3, 3:4] = T_inv

3.3 在CUDA 中对3D Gaussian 进行投影：

投影的椭圆近似成一个圆，保存圆的半径和中心点的xyz 坐标即可
当一个 Gaussian 投影成一个圆之后，那么他的半径和椭圆的标准差是密切相关的：

给定二维高斯分布的协方差矩阵 $\Sigma$ ，通过计算该矩阵的特征值并取其平方根，我们可以得到分布的“半径”，协方差的特征值可以认为是椭圆的两个方向的半径。求解下面的方程即可：

可以进一步的展开：
$(a-\lambda)(c-\lambda)-b^2=0$

圆的圆心，可以直接通过 3D Gaussian 的中心点 Center 投影得到。

计算覆盖的像素
快速的方法，将图像分成 16*16 的tile; 计算每一个 Tile 和圆的相交区域。

3.对每一个高斯按照深度顺序进行排序
一个Tile 可能会有很多个 Gaussian 进行覆盖， CUDA 程序会根据每一个 Tile 会对其覆盖的 Gaussian 进行深度排序。
计算每一个像素的颜色
每一个 Tile 对应着一个block; 而每一个Pixel 对应着一个 thread。
计算每一个像素到2D 投影圆的距离，并且依据高斯分布求解出 opacity 的大小。

Splatstudio 的后处理：

在经过一定的 Step 之后，系统会调用 Callback 函数 进行 Gaussian 的删除 Cull. 删除的频率保持和原始的 Paper 的频率一样，每隔 100 个 iterations 进行一次。当然实现的时候也有warmup=500 的设定，再500个 step 才会开始进行 Cull .

 if step % self.update_every_num_iters == 0:
                self.func(*self.args, **self.kwargs, step=step)

self.args 和 self.kwargs 都是输入的参数；要查看调用了哪个函数的名字，可以使用 self.func.__name__ 进行打印出来。

后处理包含3个部分 Densify, Cull, Reset opacity：

前 15000 个 iter 考虑的是高斯的致密化，会分裂或者复制新的 Gaussian

do_densification = (
		self.step < self.config.stop_split_at
		and self.step % reset_interval > self.num_train_data + self.config.refine_every
		)

means2d 这个变量表示的是 3D Gaussian 球投影到 2D 平面的坐标。在 CUDA 代码中，需要提前创建这个变量作为输入，然后返回值也存储在这个变量里面。

于此同时，基于 radius 这个变量，将会滤除掉不可见的 3DGS。

在 `after_train` 这个函数当中，会更新 GS Render 过程当中的一些统计量：被看见的次数，每个 GS 的梯度等

可见的 GS 指的是投影到 2D 半径 radius > 0 的高斯球，维护一个 visible_mask ：

visible_mask = (self.radii > 0).flatten()

记录每一个 GS 的在 2D 平面上的梯度 grads 和被看见的次数 vis_counts：

grads = self.xys.absgrad[0][visible_mask].norm(dim=-1)
self.vis_counts[visible_mask] += 1

记录每一个 GS 的在 2D 平面上的投影的半径占据像素平面的比例 max_2Dsize：

newradii = self.radii.detach()[visible_mask]
self.max_2Dsize[visible_mask] = torch.maximum(
	self.max_2Dsize[visible_mask],
	newradii / float(max(self.last_size[0], self.last_size[1])),
            )

大于 15000 个iter 之后，主要考虑的是高斯的删除。

 elif self.step >= self.config.stop_split_at and self.config.continue_cull_post_densification:

代码主要删除3类Gaussian：

Gaussian 的 Opacity 小于设定的 0.1

culls = (torch.sigmoid(self.opacities) < self.config.cull_alpha_thresh).squeeze()

Gaussian 的半径 Scale 过于巨大。

toobigs = (torch.exp(self.scales).max(dim=-1).values > self.config.cull_scale_thresh).squeeze()

Gaussian 投影到 2D 平面的半径占据的面积过大

toobigs =  (self.max_2Dsize > self.config.cull_screen_size).squeeze()

实际实现中，对于每一个 GS 生成一个 binary 的Mask，用来标记是否需要删除这个 GS。

在前15000 个iter 有时候需要考虑将 Gaussian 的 alpha 扩大两倍

Splatfacto 加入训练好的 checkpoint

两行代码搞定，调用父类的 load_state_dict 函数

   def load_state_dict(self, model_dict, strict=False):  # type: ignore
        super().load_state_dict(model_dict, strict= strict)

在Splatfacto 中添加 CallBack 函数进行 Zeroshot 的 Validate

目标：在 Pertrain 的阶段寻找一个最好的模型，进行 Zeroshot 的推理，得到最好的 Zeroshot 的实验结果，借助 Nerfstudio 的回调函数get_training_callbacks来完成

 def get_training_callbacks(
        self, training_callback_attributes: TrainingCallbackAttributes
    ) -> List[TrainingCallback]:
        cbs = []
        cbs.append(TrainingCallback([TrainingCallbackLocation.BEFORE_TRAIN_ITERATION], self.step_cb))
        # The order of these matters
        cbs.append(
            TrainingCallback(
                [TrainingCallbackLocation.AFTER_TRAIN_ITERATION],  ## 在训练当前 iter 之后调用
                self.after_train,  ## 具体调用的函数， 是我们的 validate 函数
                update_every_num_iters=self.config.validate_every,  ## 多久执行一次 callback 函数
                args=[training_callback_attributes.optimizers],    ## 要执行validate 函数的 输入参数
            )
        )
        return cbs

常见的 OOM Bug 记录：

经常会在 Rasterization 遇到 CUDA OOM 的问题，这个时候经常是代码出错导致的, 需要仔细的 debug 看问题出现在哪里：

render, alpha, info = rasterization(
            means=means_crop,
            quats=quats_crop / quats_crop.norm(dim=-1, keepdim=True),
            scales=scales_crop,
            opacities=opacities_crop.squeeze(-1),
            colors=colors_crop,

Solution：

NueuralSplat 里面为了更新 Gaussian 的属性值，额外添加了 update 函数，来 update 通过 MLP 学习的高斯的各种属性。如果要改成直接优化的参数，那么就不能在 update 函数里面去更新，不然会报错 OOM。

添加2个 MLP 去在200W 个分别学习 opacity 和 scale , 发现会报错 OOM

Solution：

可能还是由于 GS 属性的原因，NueuralSplat 里面的 Scale 和 Opacity 还有 feature_dc 等属性，既然都是学习而来的，那么在 GS 的 param 里面就不应该具有梯度数值，但是 torch.nn.ParameterDict 定义的参数本身会将 Dict 里面的参数 require_grad= True。将不需要优化的 Tensor 设置成 0.即可。

if self.seed_points is not None:
   self.gauss_params = torch.nn.ParameterDict(
       {
           "means": means,
           "scales": scales,
           "quats": torch.empty(0),  
           "features_dc": torch.empty(0),
           "features_rest": torch.empty(0),
           "opacities": torch.empty(0),
       }
   )