3D Gaussian Splatting主程序代码解读
文章目录
主流程
整体流程
流程图
流程简述：
1、点云初始化：初始化SfM点云，得到3D高斯球；
2、投影：借助相机外参将3D高斯点投影到图像平面上（即Splatting）；
3、光栅化：用可微光栅化，渲染得到图像；
4、计算损失：得到渲染图像Image后，将其与Ground Truth图像比较求loss，并沿蓝色箭头反向传播；
5、梯度回传：蓝色箭头向上，更新3D高斯中的参数，向下送入自适应密度控制中，更新点云。
3D高斯球
以 3D Gaussian 的方式存储信息，每一个基本单元就是一个高斯球，用一堆高斯球来表达一个场景。
每个高斯椭球都有以下变量：
中心位置 p：
三维高斯球的中心位置 p，用三维坐标（x,y,z)表示。
协方差矩阵：
R：高斯椭球的旋转矩阵，可以用四元数表达。
S：高斯椭球在各个轴的缩放矩阵，用三个实数表达。
协方差以 R和 S形式表达，总共7个参数。
体密度(透明度) α ：
一维变量。
球谐波系数：
这里使用的 J = 3（包括0，1，2，3） ，即有 16个基(系数)，那么 RGB 一共有 48 个系数。当然使用的阶数越高，模型就越精确，但是要求的系数也越多。
综上，一个高斯球总共有 59 个系数，只要给到这 59个系数，那么这个高斯球的性质就完全确定了。
球谐函数
球谐函数（Spherical Harmonics, SH）是一组定义在球面上的特殊函数，通常用来表示球面上的函数。球谐函数在图形学、计算机图形学和计算机图像等领域中广泛应用。在3DGS中，球谐函数用于近似光照和颜色分布。球谐函数的输入是theta和phi，输出是RGB颜色值。
初始化
基于SFM得到点云初始化 3D 高斯， 每个三维点初始化为一个高斯椭球。
初始化时的输入量使用的是 COLMAP 等 SFM 方式输出的点云。
3DGS 将根据这些点云进行初始化，基于这些点云的位置，会在每一个位置上放置一个高斯球，系数随机。
SfM 初始化点云
SfM (Structure from Motion，运动恢复结构) 是一种从一组图像中估计出点云的方法。
SfM 初始化点云过程的主要步骤如下：
（1）对每一张图像，使用 SIFT、SURF、ORB 等算法提取特征点，并计算特征描述子。
（2）对相邻的图像，使用 KNN、FLANN 等算法进行特征匹配，筛选出满足一致性和稳定性的匹配对。
（3）对匹配的特征点，使用 RANSAC、LMedS 等算法进行异常值剔除，提高匹配的准确性。
（4）对匹配的特征点，使用多视图几何的约束，如基础矩阵、本质矩阵、单应矩阵等，进行相机位姿的估计，以及三维坐标点的三角化。
（5）对估计的相机位姿和三维坐标点，使用 BA（Bundle Adjustment）等算法进行优化，以减少重投影误差和累积误差。
高斯初始化
3DGS 建议：从 SfM 产生的稀疏点云初始化或随机初始化高斯，可以直接调用 COLMAP 库来完成这一步，然后进行点的密集化和剪枝以控制3D高斯的密度。当由于某种原因无法获得点云时，可以使用随机初始化来代替，但可能会降低最终的重建质量。
每个点膨胀成3D高斯椭球。椭球的初始形状是一个球，使用KNN方法求3个最近邻，半径是点到3个最近邻距离的均值。
投影
给定相机位姿，将3D 高斯球投影到图像上。即：3D 高斯（椭球）被投影到 2D 图像空间（椭圆）中进行渲染。投影依据的公式如下：
高斯点均值（随机变量x的高斯概率密度函数）：
3D协方差矩阵：
给定视图变换W 和 3D 协方差矩阵 ∑，雅可比矩阵J，投影的2D协方差矩阵：
投影的 2D 高斯的中心位置和颜色可以直接从 3D 高斯的参数中得到。投影的 2D 高斯的不透明度需要根据 3D 高斯的不透明度和协方差矩阵进行调整。投影的 2D 高斯的不透明度，可以使用以下公式计算：
可微光栅化
3DGS 从近到远每个球投下来以后都形成了一个图像区，那么在重叠区域就可以进行光栅化的融合了。每个点都进行融合以后就可以得到图像。此环节涉及分块、排序、α-blending这几个步骤。
给定像素位置x，通过视图变换 W，可以计算与所有重叠高斯体的距离，即这些高斯体的深度，形成高斯体的排序列表 N。然后，进行Alpha Blending，也就是混合 alpha 合成来计算整体图像的最终颜色。
计算损失
损失计算公式
L1：度量两像素间差异；
LD-SSIM：度量两图像间结果差异。
λ：优化器参数（系统默认为0.02）。
在 3DGS 内，每次采集一小批图，以图像为单位进行损失计算。
梯度回传
更新高斯属性
更新3D 高斯球的59 个属性。
自适应密度控制
自适应高斯稠密化方案：
学习过程中，较大梯度的高斯椭球存在 欠重构(under-reconstruction) 和 过重构(over-reconstruction) 问题。
梯度在传过来时没有更新任何参数，只是通过对这 59维导数的模值来确定当前高斯球，是否存在欠重构或过重构的问题，如果是就进行复制或分裂。这个步骤是不可导的。
① 欠重构区域的高斯椭球方差小， 进行复制操作；
可以看到上图中的几何体，又是很难用一个高斯球去描述这个几何体的形状，所以就对高斯球进行克隆，克隆的操作是不可导的。克隆完再优化就成了右边的样子。
② 过重构区域的高斯椭球方差大， 进行分裂操作；
图中可以看到方差大的高斯球太大了，拟合覆盖了全部形状，但有太多不属于这个几何形体的形状，这样描述是不准确的。
③ 每经过固定次数的迭代进行一次剔除操作， 剔除几乎透明的高斯椭球以及方差过大的高斯椭球。
伪代码
核心代码解读
代码路径：train.py
主函数
if __name__ == "__main__":
# Set up command line argument parser
parser = ArgumentParser(description="Training script parameters")
lp = ModelParams(parser)
op = OptimizationParams(parser)
pp = PipelineParams(parser)
parser.add_argument('--ip', type=str, default="127.0.0.1")
parser.add_argument('--port', type=int, default=6009)
parser.add_argument('--debug_from', type=int, default=-1)
parser.add_argument('--detect_anomaly', action='store_true', default=False)
parser.add_argument("--test_iterations", nargs="+", type=int, default=[7_000, 30_000])
parser.add_argument("--save_iterations", nargs="+", type=int, default=[7_000, 30_000])
parser.add_argument("--quiet", action="store_true")
parser.add_argument("--checkpoint_iterations", nargs="+", type=int, default=[])
parser.add_argument("--start_checkpoint", type=str, default = None)
args = parser.parse_args(sys.argv[1:])
args.save_iterations.append(args.iterations)
print("Optimizing " + args.model_path)
# Initialize system state (RNG)
safe_state(args.quiet)
# Start GUI server, configure and run training
#这行代码初始化一个 GUI 服务器，使用 args.ip 和 args.port 作为参数。这可能是一个用于监视和控制训练过程的图形用户界面的一部分。
network_gui.init(args.ip, args.port)
#这行代码设置 PyTorch 是否要检测梯度计算中的异常。
torch.autograd.set_detect_anomaly(args.detect_anomaly)
# 输入的参数包括：模型的参数（传入的为数据集的位置）、优化器的参数、其他pipeline的参数，测试迭代次数、保存迭代次数 、检查点迭代次数 、开始检查点 、调试起点
training(lp.extract(args), op.extract(args), pp.extract(args), args.test_iterations, args.save_iterations, args.checkpoint_iterations, args.start_checkpoint, args.debug_from)
# All done
print("\nTraining complete.")
初始化系统状态
def safe_state(silent):
old_f = sys.stdout
class F:
def __init__(self, silent):
self.silent = silent
def write(self, x):
if not self.silent:
if x.endswith("\n"):
old_f.write(x.replace("\n", " [{}]\n".format(str(datetime.now().strftime("%d/%m %H:%M:%S")))))
else:
old_f.write(x)
def flush(self):
old_f.flush()
sys.stdout = F(silent)
random.seed(0)
np.random.seed(0)
torch.manual_seed(0)
torch.cuda.set_device(torch.device("cuda:0"))
GUI初始化、PyTorch设置，此处不展开讲解。
模型训练
训练的流程
训练过程每个迭代主要执行以下的操作：
- 每 1000 次迭代，增加球谐系数的阶数。
- 随机选择一个相机视角。
- 渲染图像，获取视点空间点、能见度过滤器和半径等信息。
- 计算损失（L1 损失和 DSSIM 损失的加权和），进行反向传播。
- 通过无梯度的上下文进行后续操作：
– 根据迭代次数进行点云密度操作（densification）。
– 更新最大半径信息。
– 根据条件进行点云密度增加和修剪。
– 进行优化器的参数更新。
在整个训练过程中，这些步骤循环执行，逐渐优化模型参数，进行损失计算和反向传播，同时根据条件进行点云密度操作和保存检查点，以逐步提升模型性能。
training函数
def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoint_iterations, checkpoint, debug_from):
#初始化迭代次数。#初始化迭代次数。
first_iter = 0
#设置 TensorBoard 写入器和日志记录器。
tb_writer = prepare_output_and_logger(dataset)
创建高斯模型和场景，设置训练参数
#（重点看，需要转跳）创建一个 GaussianModel 类的实例，输入一系列参数，其参数取自数据集。
gaussians = GaussianModel(dataset.sh_degree)
#（这个类的主要目的是处理场景的初始化、保存和获取相机信息等任务，）创建一个 Scene 类的实例，使用数据集和之前创建的 GaussianModel 实例作为参数。
scene = Scene(dataset, gaussians)
#设置 GaussianModel 的训练参数。
gaussians.training_setup(opt)
模型训练
#如果有提供检查点路径。
if checkpoint:
#通过 torch.load(checkpoint) 加载检查点的模型参数和起始迭代次数。
(model_params, first_iter) = torch.load(checkpoint)
#通过 gaussians.restore 恢复模型的状态。
gaussians.restore(model_params, opt)
#设置背景颜色，根据数据集是否有白色背景来选择。
bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]
#将背景颜色转化为 PyTorch Tensor，并移到 GPU 上。
background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
# 创建两个 CUDA 事件，用于测量迭代时间。
iter_start = torch.cuda.Event(enable_timing = True)
iter_end = torch.cuda.Event(enable_timing = True)
viewpoint_stack = None
ema_loss_for_log = 0.0
#创建一个 tqdm 进度条，用于显示训练进度。
progress_bar = tqdm(range(first_iter, opt.iterations), desc="Training progress")
模型训练循环迭代
first_iter += 1
# 接下来开始循环迭代
#主要的训练循环开始。
for iteration in range(first_iter, opt.iterations + 1):
#检查 GUI 是否连接，如果连接则接收 GUI 发送的消息。
if network_gui.conn == None:
network_gui.try_connect()
while network_gui.conn != None:
try:
net_image_bytes = None
custom_cam, do_training, pipe.convert_SHs_python, pipe.compute_cov3D_python, keep_alive, scaling_modifer = network_gui.receive()
if custom_cam != None:
net_image = render(custom_cam, gaussians, pipe, background, scaling_modifer)["render"]
net_image_bytes = memoryview((torch.clamp(net_image, min=0, max=1.0) * 255).byte().permute(1, 2, 0).contiguous().cpu().numpy())
network_gui.send(net_image_bytes, dataset.source_path)
if do_training and ((iteration < int(opt.iterations)) or not keep_alive):
break
except Exception as e:
network_gui.conn = None
#用于测量迭代时间。
iter_start.record()
#更新学习率。
gaussians.update_learning_rate(iteration)
增加球谐函数的阶数
# Every 1000 its we increase the levels of SH up to a maximum degree
#每 1000 次迭代，增加球谐函数的阶数。
if iteration % 1000 == 0:
gaussians.oneupSHdegree()
随机选择一个相机视角
# Pick a random Camera （随机选择一个训练相机。）
if not viewpoint_stack:
viewpoint_stack = scene.getTrainCameras().copy()
viewpoint_cam = viewpoint_stack.pop(randint(0, len(viewpoint_stack)-1))
渲染图像
# Render （渲染图像，计算损失（L1 loss 和 SSIM loss））
if (iteration - 1) == debug_from:
pipe.debug = True
bg = torch.rand((3), device="cuda") if opt.random_background else background
render_pkg = render(viewpoint_cam, gaussians, pipe, bg)
image, viewspace_point_tensor, visibility_filter, radii = render_pkg["render"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["radii"]
计算损失loss
# Loss
gt_image = viewpoint_cam.original_image.cuda()
Ll1 = l1_loss(image, gt_image)
#计算渲染的图像与真实图像之间的loss
loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim(image, gt_image))
loss反向传播
#更新损失。loss反向传播
loss.backward()
#用于测量迭代时间。
iter_end.record()
#记录损失的指数移动平均值，并定期更新进度条。
with torch.no_grad():
# Progress bar
ema_loss_for_log = 0.4 * loss.item() + 0.6 * ema_loss_for_log
if iteration % 10 == 0:
progress_bar.set_postfix({"Loss": f"{ema_loss_for_log:.{7}f}"})
progress_bar.update(10)
if iteration == opt.iterations:
progress_bar.close()
日志和场景保存
# Log and save
training_report(tb_writer, iteration, Ll1, loss, l1_loss, iter_start.elapsed_time(iter_end), testing_iterations, scene, render, (pipe, background))
#如果达到保存迭代次数，保存场景。
if (iteration in saving_iterations):
print("\n[ITER {}] Saving Gaussians".format(iteration))
scene.save(iteration)
密集化处理
# Densification（在一定的迭代次数内进行密集化处理。）
#在达到指定的迭代次数之前执行以下操作。
if iteration < opt.densify_until_iter:
# Keep track of max radii in image-space for pruning
#将每个像素位置上的最大半径记录在 max_radii2D 中。这是为了密集化时进行修剪（pruning）操作时的参考。
gaussians.max_radii2D[visibility_filter] = torch.max(gaussians.max_radii2D[visibility_filter], radii[visibility_filter])
#将与密集化相关的统计信息添加到 gaussians 模型中，包括视图空间点和可见性过滤器。
gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter)
#在指定的迭代次数之后，每隔一定的迭代间隔进行以下密集化操作。
if iteration > opt.densify_from_iter and iteration % opt.densification_interval == 0:
#根据当前迭代次数设置密集化的阈值。如果当前迭代次数大于 opt.opacity_reset_interval，则设置 size_threshold 为 20，否则为 None。
size_threshold = 20 if iteration > opt.opacity_reset_interval else None
#执行密集化和修剪操作，其中包括梯度阈值、密集化阈值、相机范围和之前计算的 size_threshold。
gaussians.densify_and_prune(opt.densify_grad_threshold, 0.005, scene.cameras_extent, size_threshold)
#在每隔一定迭代次数或在白色背景数据集上的指定迭代次数时，执行以下操作。
if iteration % opt.opacity_reset_interval == 0 or (dataset.white_background and iteration == opt.densify_from_iter):
#重置模型中的某些参数，涉及到透明度的操作，具体实现可以在 reset_opacity 方法中找到。
gaussians.reset_opacity()
执行优化器
# Optimizer step（执行优化器的步骤，然后清零梯度。）
if iteration < opt.iterations:
gaussians.optimizer.step()
gaussians.optimizer.zero_grad(set_to_none = True)
保存检查点
# 如果达到检查点迭代次数，保存检查点。
if (iteration in checkpoint_iterations):
print("\n[ITER {}] Saving Checkpoint".format(iteration))
torch.save((gaussians.capture(), iteration), scene.model_path + "/chkpnt" + str(iteration) + ".pth")
场景初始化与操作
场景类，用于管理和加载 3D 场景的参数、模型和相机信息，并支持不同分辨率的相机数据。Scene 类结合了高斯模型和数据集处理逻辑，尤其适用于从 COLMAP 或 Blender 数据集中加载相机和场景信息，以便进行 3D 表示和训练。
代码路径：scene_init_.py
class Scene:
gaussians : GaussianModel
# 初始化方法
def __init__(self, args : ModelParams, gaussians : GaussianModel, load_iteration=None, shuffle=True, resolution_scales=[1.0]):
"""b
:param path: Path to colmap scene main folder.
"""
self.model_path = args.model_path #将传入的 args 对象中的 model_path 属性赋值给 self.model_path，表示模型的路径。
self.loaded_iter = None #初始化 self.loaded_iter 为 None，用于存储已加载的模型的迭代次数。
self.gaussians = gaussians #将传入的高斯模型对象赋值给 self.gaussians 属性。
if load_iteration: #可选参数，默认为 None。如果提供了值，它将被用作已加载模型的迭代次数。
if load_iteration == -1: #如果没有提供 load_iteration，则将点云数据和相机信息保存到文件中。
self.loaded_iter = searchForMaxIteration(os.path.join(self.model_path, "point_cloud"))
else:
self.loaded_iter = load_iteration
print("Loading trained model at iteration {}".format(self.loaded_iter)) #输出加载模型的迭代次数的信息。
self.train_cameras = {}
self.test_cameras = {}
# 根据场景的类型（Colmap 或 Blender）加载相应的场景信息，存储在 scene_info 变量中。
if os.path.exists(os.path.join(args.source_path, "sparse")):
scene_info = sceneLoadTypeCallbacks["Colmap"](args.source_path, args.images, args.eval)
elif os.path.exists(os.path.join(args.source_path, "transforms_train.json")):
print("Found transforms_train.json file, assuming Blender data set!")
scene_info = sceneLoadTypeCallbacks["Blender"](args.source_path, args.white_background, args.eval)
else:
assert False, "Could not recognize scene type!"
# 保存点云数据和相机信息：
if not self.loaded_iter:
with open(scene_info.ply_path, 'rb') as src_file, open(os.path.join(self.model_path, "input.ply") , 'wb') as dest_file:
dest_file.write(src_file.read())
json_cams = []
camlist = []
if scene_info.test_cameras:
camlist.extend(scene_info.test_cameras)
if scene_info.train_cameras:
camlist.extend(scene_info.train_cameras)
for id, cam in enumerate(camlist):
json_cams.append(camera_to_JSON(id, cam))
with open(os.path.join(self.model_path, "cameras.json"), 'w') as file:
json.dump(json_cams, file)
# 随机排序相机
if shuffle: #可选参数，默认为 True。如果设置为 True，则会对场景中的训练和测试相机进行随机排序。
random.shuffle(scene_info.train_cameras) # Multi-res consistent random shuffling
random.shuffle(scene_info.test_cameras) # Multi-res consistent random shuffling
# 设置相机的范围：
self.cameras_extent = scene_info.nerf_normalization["radius"]
# 加载训练和测试相机：
for resolution_scale in resolution_scales: #可选参数，默认为 [1.0]。一个浮点数列表，用于指定训练和测试相机的分辨率缩放因子。
print("Loading Training Cameras")
self.train_cameras[resolution_scale] = cameraList_from_camInfos(scene_info.train_cameras, resolution_scale, args)
print("Loading Test Cameras")
self.test_cameras[resolution_scale] = cameraList_from_camInfos(scene_info.test_cameras, resolution_scale, args)
# 加载或创建高斯模型
if self.loaded_iter: #如果已加载模型，则调用 load_ply 方法加载点云数据。
self.gaussians.load_ply(os.path.join(self.model_path,
"point_cloud",
"iteration_" + str(self.loaded_iter),
"point_cloud.ply"))
else: #否则，调用 create_from_pcd 方法根据场景信息中的点云数据创建高斯模型。
self.gaussians.create_from_pcd(scene_info.point_cloud, self.cameras_extent)
def save(self, iteration):#该方法用于保存点云数据到文件，其参数 iteration 指定了迭代次数。
point_cloud_path = os.path.join(self.model_path, "point_cloud/iteration_{}".format(iteration))
self.gaussians.save_ply(os.path.join(point_cloud_path, "point_cloud.ply"))
def getTrainCameras(self, scale=1.0):#返回训练相机的列表，可以根据指定的缩放因子 scale 获取相应分辨率的相机列表。
return self.train_cameras[scale]
def getTestCameras(self, scale=1.0):#返回测试相机的列表，可以根据指定的缩放因子 scale 获取相应分辨率的相机列表。
return self.test_cameras[scale]
根据场景类型加载场景信息，具体实现请查看scene\dataset_readers.py。
高斯模型初始化和加载
scene_init_.py
if self.loaded_iter: #如果已加载模型，则调用 load_ply 方法加载点云数据。
self.gaussians.load_ply(os.path.join(self.model_path,
"point_cloud",
"iteration_" + str(self.loaded_iter),
"point_cloud.ply"))
else: #否则，调用 create_from_pcd 方法根据场景信息中的点云数据创建高斯模型。
self.gaussians.create_from_pcd(scene_info.point_cloud, self.cameras_extent)
高斯模型定义
初始化与高斯模型相关的各种属性，如使用的球谐阶数、最大球谐阶数、各种张量（_xyz、_features_dc等）、优化器和其他参数。
调用 setup_functions 方法设置各种激活和变换函数。
代码路径：scene\gaussian_model.py
def __init__(self, sh_degree : int):
self.active_sh_degree = 0 #球谐阶数
self.max_sh_degree = sh_degree #最大球谐阶数
# 存储不同信息的张量（tensor）
self._xyz = torch.empty(0) #空间位置
self._features_dc = torch.empty(0)
self._features_rest = torch.empty(0)
self._scaling = torch.empty(0) #椭球的形状尺度
self._rotation = torch.empty(0) #椭球的旋转
self._opacity = torch.empty(0) #不透明度
self.max_radii2D = torch.empty(0)
self.xyz_gradient_accum = torch.empty(0)
self.denom = torch.empty(0)
self.optimizer = None #初始化优化器为 None。
self.percent_dense = 0 #初始化百分比密度为0。
self.spatial_lr_scale = 0 #初始化空间学习速率缩放为0。
self.setup_functions() #调用 setup_functions 方法设置各种激活和变换函数
def setup_functions(self): #用于设置一些激活函数和变换函数
def build_covariance_from_scaling_rotation(scaling, scaling_modifier, 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 #将尺度激活函数设置为指数函数。
self.scaling_inverse_activation = torch.log #将尺度逆激活函数设置为对数函数。
self.covariance_activation = build_covariance_from_scaling_rotation #将协方差激活函数设置为上述定义的 build_covariance_from_scaling_rotation 函数。
self.opacity_activation = torch.sigmoid #将不透明度激活函数设置为 sigmoid 函数。
self.inverse_opacity_activation = inverse_sigmoid #将不透明度逆激活函数设置为一个名为 inverse_sigmoid 的函数
self.rotation_activation = torch.nn.functional.normalize #用于归一化旋转矩阵。
高斯模型创建
scene\gaussian_model.py
def create_from_pcd(self, pcd : BasicPointCloud, spatial_lr_scale : float): #用于从给定的点云数据 pcd 创建对象的初始化状态。
self.spatial_lr_scale = spatial_lr_scale
fused_point_cloud = torch.tensor(np.asarray(pcd.points)).float().cuda()
fused_color = RGB2SH(torch.tensor(np.asarray(pcd.colors)).float().cuda())
features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda()
features[:, :3, 0 ] = fused_color
features[:, 3:, 1:] = 0.0
print("Number of points at initialisation : ", fused_point_cloud.shape[0])
dist2 = torch.clamp_min(distCUDA2(torch.from_numpy(np.asarray(pcd.points)).float().cuda()), 0.0000001)
scales = torch.log(torch.sqrt(dist2))[...,None].repeat(1, 3)
rots = torch.zeros((fused_point_cloud.shape[0], 4), device="cuda")
rots[:, 0] = 1
opacities = inverse_sigmoid(0.1 * torch.ones((fused_point_cloud.shape[0], 1), dtype=torch.float, device="cuda"))
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))
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")
distCUDA2的定义
submodules\simple-knn\ext.cpp
#include <torch/extension.h>
#include "spatial.h"
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("distCUDA2", &distCUDA2);
}
distCUDA2的实现
详见代码文件：submodules\simple-knn\spatial.cu
高斯模型加载
scene\gaussian_model.py
def load_ply(self, path): #这个方法的目的是从PLY文件中加载各种数据，并将这些数据存储为类中的属性，以便后续的操作和训练。
plydata = PlyData.read(path)
xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
np.asarray(plydata.elements[0]["y"]),
np.asarray(plydata.elements[0]["z"])), axis=1)
opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]
features_dc = np.zeros((xyz.shape[0], 3, 1))
features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"])
features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"])
extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")]
extra_f_names = sorted(extra_f_names, key = lambda x: int(x.split('_')[-1]))
assert len(extra_f_names)==3*(self.max_sh_degree + 1) ** 2 - 3
features_extra = np.zeros((xyz.shape[0], len(extra_f_names)))
for idx, attr_name in enumerate(extra_f_names):
features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name])
# Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC)
features_extra = features_extra.reshape((features_extra.shape[0], 3, (self.max_sh_degree + 1) ** 2 - 1))
scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
scale_names = sorted(scale_names, key = lambda x: int(x.split('_')[-1]))
scales = np.zeros((xyz.shape[0], len(scale_names)))
for idx, attr_name in enumerate(scale_names):
scales[:, idx] = np.asarray(plydata.elements[0][attr_name])
rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")]
rot_names = sorted(rot_names, key = lambda x: int(x.split('_')[-1]))
rots = np.zeros((xyz.shape[0], len(rot_names)))
for idx, attr_name in enumerate(rot_names):
rots[:, idx] = np.asarray(plydata.elements[0][attr_name])
self._xyz = nn.Parameter(torch.tensor(xyz, dtype=torch.float, device="cuda").requires_grad_(True))
self._features_dc = nn.Parameter(torch.tensor(features_dc, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True))
self._features_rest = nn.Parameter(torch.tensor(features_extra, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True))
self._opacity = nn.Parameter(torch.tensor(opacities, dtype=torch.float, device="cuda").requires_grad_(True))
self._scaling = nn.Parameter(torch.tensor(scales, dtype=torch.float, device="cuda").requires_grad_(True))
self._rotation = nn.Parameter(torch.tensor(rots, dtype=torch.float, device="cuda").requires_grad_(True))
self.active_sh_degree = self.max_sh_degree
自适应密集化处理
点密集化：在点密集化阶段，3DGS自适应地增加高斯的密度，以更好地捕捉场景的细节。该过程特别关注缺失几何特征或高斯过于分散的区域。密集化在一定数量的迭代后执行，比如100个迭代，针对在视图空间中具有较大位置梯度（即超过特定阈值）的高斯。其包括在未充分重建的区域克隆小高斯或在过度重建的区域分裂大高斯。对于克隆，创建高斯的复制体并朝着位置梯度移动。对于分裂，用两个较小的高斯替换一个大高斯，按照特定因子减小它们的尺度。这一步旨在于 3D 空间中寻求高斯的最佳分布和表示，增强重建的整体质量。
点的剪枝：点的剪枝阶段移除冗余或影响较小的高斯，可以在某种程度上看作是一种正则化过程。一般消除几乎是透明的高斯（α低于指定阈值）和在世界空间或视图空间中过大的高斯。此外，为防止输入相机附近的高斯密度不合理地增加，这些高斯会在固定次数的迭代后，将
设置为接近0的值。该步骤在保证高斯的精度和有效性的情况下，能节约计算资源。
def densify_and_prune(self, max_grad, min_opacity, extent, max_screen_size):
grads = self.xyz_gradient_accum / self.denom #计算密度估计的梯度
grads[grads.isnan()] = 0.0 #将梯度中的 NaN（非数值）值设置为零，以处理可能的数值不稳定性。
self.densify_and_clone(grads, max_grad, extent) #对under reconstruction的区域进行稠密化和复制操作
self.densify_and_split(grads, max_grad, extent) #对over reconstruction的区域进行稠密化和分割操作
prune_mask = (self.get_opacity < min_opacity).squeeze() #创建一个掩码，标记那些透明度小于指定阈值的点。.squeeze() 用于去除掩码中的单维度。
if max_screen_size: #如何设置了相机的范围，
big_points_vs = self.max_radii2D > max_screen_size #创建一个掩码，标记在图像空间中半径大于指定阈值的点。
big_points_ws = self.get_scaling.max(dim=1).values > 0.1 * extent #创建一个掩码，标记在世界空间中尺寸大于指定阈值的点。
prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws) #将这两个掩码与先前的透明度掩码进行逻辑或操作，得到最终的修剪掩码。
self.prune_points(prune_mask) #：根据修剪掩码，修剪模型中的一些参数。
torch.cuda.empty_cache() #清理 GPU 缓存，释放一些内存
克隆
densify_and_clone 函数
和 densify_and_split 函数都是用于点云稠密化的方法，但它们的策略和行为有所不同：
目的：在满足梯度条件的基础上，将满足缩放和不透明度条件的点进行复制，以增加点云的密度。
行为：对于满足梯度条件、缩放条件和不透明度条件的点，会将它们复制多次，以形成新的点集，从而增加了点的数量。
策略：使用原始点的属性，如坐标、缩放、旋转、特征和不透明度，创建新的点集。
def densify_and_clone(self, grads, grad_threshold, scene_extent):
# Extract points that satisfy the gradient condition
selected_pts_mask = torch.where(torch.norm(grads, dim=-1) >= grad_threshold, True, False) #建一个掩码，标记满足梯度条件的点。具体来说，对于每个点，计算其梯度的L2范数，如果大于等于指定的梯度阈值，则标记为True，否则标记为False。
selected_pts_mask = torch.logical_and(selected_pts_mask,
torch.max(self.get_scaling, dim=1).values <= self.percent_dense*scene_extent)
# 在上述掩码的基础上，进一步过滤掉那些缩放（scaling）大于一定百分比（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)
修剪
densify_and_split
这个函数的主要目的是在满足梯度条件的情况下对点进行稠密化和分割操作，以增加点的数量并改善点云的表示。函数步骤如下：
获取初始点数 n_init_points。
创建布尔掩码 selected_pts_mask，将满足梯度条件且缩放值大于 percent_dense * scene_extent 的点选中。
为满足条件的点创建分割坐标和属性，并将它们重复 N 次，形成 N 个分割点集。
计算分割点的缩放值、旋转值和其他属性。
调用 densification_postfix 方法，对新的点和属性进行后处理。
创建布尔掩码 prune_filter，其中包括选中的点和新生成的分割点。
使用 prune_points 方法进行点的修剪。
def densify_and_split(self, grads, grad_threshold, scene_extent, N=2):
n_init_points = self.get_xyz.shape[0] #获取初始点的数量。
# Extract points that satisfy the gradient condition
padded_grad = torch.zeros((n_init_points), device="cuda") #创建一个长度为初始点数量的梯度张量，并将计算得到的梯度填充到其中。
padded_grad[:grads.shape[0]] = grads.squeeze()
selected_pts_mask = torch.where(padded_grad >= grad_threshold, True, False) #创建一个掩码，标记那些梯度大于等于指定阈值的点。
selected_pts_mask = torch.logical_and(selected_pts_mask,
torch.max(self.get_scaling, dim=1).values > self.percent_dense*scene_extent)
# 一步过滤掉那些缩放（scaling）大于一定百分比的场景范围的点。
# 为每个点生成新的样本，其中 stds 是点的缩放，means 是均值。
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) #为每个点构建旋转矩阵，并将其重复 N 次。
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) #将旋转矩阵重复 N 次。
# 将原始点的特征重复 N 次。
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)
# 调用另一个方法 densification_postfix，该方法对新生成的点执行后处理操作（此处跟densify_and_clone一样）。
self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacity, new_scaling, new_rotation)
# 创建一个修剪（pruning）的过滤器，将新生成的点添加到原始点的掩码之后。
prune_filter = torch.cat((selected_pts_mask, torch.zeros(N * selected_pts_mask.sum(), device="cuda", dtype=bool)))
# 根据修剪过滤器，修剪模型中的一些参数。
self.prune_points(prune_filter)
build_rotation的具体实现在utils\general_utils.py中。
densification_postfix的具体实现在当前文件scene\gaussian_model.py。
densification_postfix：将新的密集化点的相关特征保存在一个字典中。
def densification_postfix(self, new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation):
d = {"xyz": new_xyz,
"f_dc": new_features_dc,
"f_rest": new_features_rest,
"opacity": new_opacities,
"scaling" : new_scaling,
"rotation" : new_rotation}
#将字典中的张量连接（concatenate）成可优化的张量。这个方法的具体实现可能是将字典中的每个张量进行堆叠，以便于在优化器中进行处理。
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")
张量优化
def cat_tensors_to_optimizer(self, tensors_dict):
optimizable_tensors = {}
for group in self.optimizer.param_groups:
assert len(group["params"]) == 1
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_state
optimizable_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
prune_points
该函数根据给定的布尔掩码，从点云中移除那些需要剪枝的点。
该方法会调用 _prune_optimizer 方法，将有效点的掩码传递给优化器，以从优化器中剪枝相关的参数。
最后更新实例变量 self._xyz、self._features_dc、self._features_rest、self._opacity、self._scaling 和 self._rotation，以保留有效点的相关属性。
def prune_points(self, mask):
valid_points_mask = ~mask
optimizable_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]
_prune_optimizer
def _prune_optimizer(self, mask):
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_state
optimizable_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
渲染及光栅化
render的实现在gaussian_renderer_init_.py中。
场景渲染
这段代码是一个用于渲染场景的函数，主要是通过将高斯分布的点投影到2D屏幕上来生成渲染图像。
def render(viewpoint_camera, pc : GaussianModel, pipe, bg_color : torch.Tensor, scaling_modifier = 1.0, override_color = None):
"""
Render the scene.
Background tensor (bg_color) must be on GPU!
"""
# Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
# 创建一个与输入点云（高斯模型）大小相同的零张量，用于记录屏幕空间中的点的位置。这个张量将用于计算对于屏幕空间坐标的梯度。
screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
try:
screenspace_points.retain_grad() #尝试保留张量的梯度。这是为了确保可以在反向传播过程中计算对于屏幕空间坐标的梯度。
except:
pass
光栅化设置
# Set up rasterization configuration
# 计算视场的 tan 值，这将用于设置光栅化配置。
tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)
# 设置光栅化的配置，包括图像的大小、视场的 tan 值、背景颜色、视图矩阵、投影矩阵等。
raster_settings = GaussianRasterizationSettings(
image_height=int(viewpoint_camera.image_height),
image_width=int(viewpoint_camera.image_width),
tanfovx=tanfovx,
tanfovy=tanfovy,
bg=bg_color,
scale_modifier=scaling_modifier,
viewmatrix=viewpoint_camera.world_view_transform,
projmatrix=viewpoint_camera.full_proj_transform,
sh_degree=pc.active_sh_degree,
campos=viewpoint_camera.camera_center,
prefiltered=False,
debug=pipe.debug
)
创建高斯光栅化器对象
#创建一个高斯光栅化器对象，用于将高斯分布投影到屏幕上。
rasterizer = GaussianRasterizer(raster_settings=raster_settings)
# 获取高斯分布的三维坐标、屏幕空间坐标和透明度。
means3D = pc.get_xyz
means2D = screenspace_points
opacity = pc.get_opacity
# If precomputed 3d covariance is provided, use it. If not, then it will be computed from
# scaling / rotation by the rasterizer.
# 如果提供了预先计算的3D协方差矩阵，则使用它。否则，它将由光栅化器根据尺度和旋转进行计算。
scales = None
rotations = None
cov3D_precomp = None
if pipe.compute_cov3D_python:
cov3D_precomp = pc.get_covariance(scaling_modifier) #获取预计算的三维协方差矩阵。
else: #获取缩放和旋转信息。（对应的就是3D高斯的协方差矩阵了）
scales = pc.get_scaling
rotations = pc.get_rotation
# If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
# from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
# 如果提供了预先计算的颜色，则使用它们。否则，如果希望在Python中从球谐函数中预计算颜色，请执行此操作。如果没有，则颜色将通过光栅化器进行从球谐函数到RGB的转换。
shs = None
colors_precomp = None
if override_color is None:
if pipe.convert_SHs_python:
shs_view = pc.get_features.transpose(1, 2).view(-1, 3, (pc.max_sh_degree+1)**2) #将SH特征的形状调整为（batch_size * num_points，3，(max_sh_degree+1)**2）。
dir_pp = (pc.get_xyz - viewpoint_camera.camera_center.repeat(pc.get_features.shape[0], 1)) #计算相机中心到每个点的方向向量，并归一化。
dir_pp_normalized = dir_pp/dir_pp.norm(dim=1, keepdim=True) #计算相机中心到每个点的方向向量，并归一化。
sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized) #使用SH特征将方向向量转换为RGB颜色。
colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0) #将RGB颜色的范围限制在0到1之间。
else:
shs = pc.get_features
else:
colors_precomp = override_color
调用光栅化器
# Rasterize visible Gaussians to image, obtain their radii (on screen).
# 调用光栅化器，将高斯分布投影到屏幕上，获得渲染图像和每个高斯分布在屏幕上的半径。
rendered_image, radii = rasterizer(
means3D = means3D,
means2D = means2D,
shs = shs,
colors_precomp = colors_precomp,
opacities = opacity,
scales = scales,
rotations = rotations,
cov3D_precomp = cov3D_precomp)
# Those Gaussians that were frustum culled or had a radius of 0 were not visible.
# They will be excluded from value updates used in the splitting criteria.
# 返回一个字典，包含渲染的图像、屏幕空间坐标、可见性过滤器（根据半径判断是否可见）以及每个高斯分布在屏幕上的半径。
return {"render": rendered_image,
"viewspace_points": screenspace_points,
"visibility_filter" : radii > 0,
"radii": radii}
高斯光栅化器
代码路径：submodules\diff-gaussian-rasterization\diff_gaussian_rasterization_init_.py
高斯光栅化PyTorch Autograd Function
这是一个自定义的 PyTorch Autograd Function，用于高斯光栅化的前向传播和反向传播。
class _RasterizeGaussians(torch.autograd.Function):
@staticmethod
def forward( #用于定义前向渲染的规则，接受一系列输入参数，并调用 C++/CUDA 实现的 _C.rasterize_gaussians 方法进行高斯光栅化。
ctx, #上下文对象，用于保存计算中间结果以供反向传播使用。（后面几个是输入参数。）
means3D,
means2D,
sh,
colors_precomp,
opacities,
scales,
rotations,
cov3Ds_precomp,
raster_settings,
):
# Restructure arguments the way that the C++ lib expects them
args = (
raster_settings.bg,
means3D,
colors_precomp,
opacities,
scales,
rotations,
raster_settings.scale_modifier,
cov3Ds_precomp,
raster_settings.viewmatrix,
raster_settings.projmatrix,
raster_settings.tanfovx,
raster_settings.tanfovy,
raster_settings.image_height,
raster_settings.image_width,
sh,
raster_settings.sh_degree,
raster_settings.campos,
raster_settings.prefiltered,
raster_settings.debug
)
# Invoke C++/CUDA rasterizer
if raster_settings.debug:
cpu_args = cpu_deep_copy_tuple(args) # Copy them before they can be corrupted
try:
num_rendered, color, radii, geomBuffer, binningBuffer, imgBuffer = _C.rasterize_gaussians(*args) #C++/CUDA 光栅化计算的输出结果。
except Exception as ex:
torch.save(cpu_args, "snapshot_fw.dump")
print("\nAn error occured in forward. Please forward snapshot_fw.dump for debugging.")
raise ex
else:
num_rendered, color, radii, geomBuffer, binningBuffer, imgBuffer = _C.rasterize_gaussians(*args)
# Keep relevant tensors for backward
ctx.raster_settings = raster_settings
ctx.num_rendered = num_rendered
ctx.save_for_backward(colors_precomp, means3D, scales, rotations, cov3Ds_precomp, radii, sh, geomBuffer, binningBuffer, imgBuffer)
return color, radii
@staticmethod
def backward(ctx, grad_out_color, _): #方法用于定义反向传播梯度下降的规则，接受输入的梯度
# Restore necessary values from context
num_rendered = ctx.num_rendered
raster_settings = ctx.raster_settings
colors_precomp, means3D, scales, rotations, cov3Ds_precomp, radii, sh, geomBuffer, binningBuffer, imgBuffer = ctx.saved_tensors
# Restructure args as C++ method expects them
# 将梯度和其他输入参数重构为 C++ 方法所期望的形式。
args = (raster_settings.bg,
means3D,
radii,
colors_precomp,
scales,
rotations,
raster_settings.scale_modifier,
cov3Ds_precomp,
raster_settings.viewmatrix,
raster_settings.projmatrix,
raster_settings.tanfovx,
raster_settings.tanfovy,
grad_out_color,
sh,
raster_settings.sh_degree,
raster_settings.campos,
geomBuffer,
num_rendered,
binningBuffer,
imgBuffer,
raster_settings.debug)
# Compute gradients for relevant tensors by invoking backward method
# 注意，该函数中包含了对调试模式的处理，即如果启用了调试模式，则在计算前向和反向传播时保存了参数的副本，并在出现异常时将其保存到文件中，以供调试。
if raster_settings.debug:
cpu_args = cpu_deep_copy_tuple(args) # Copy them before they can be corrupted
try:
grad_means2D, grad_colors_precomp, grad_opacities, grad_means3D, grad_cov3Ds_precomp, grad_sh, grad_scales, grad_rotations = _C.rasterize_gaussians_backward(*args)
except Exception as ex:
torch.save(cpu_args, "snapshot_bw.dump")
print("\nAn error occured in backward. Writing snapshot_bw.dump for debugging.\n")
raise ex
else:
grad_means2D, grad_colors_precomp, grad_opacities, grad_means3D, grad_cov3Ds_precomp, grad_sh, grad_scales, grad_rotations = _C.rasterize_gaussians_backward(*args)
#梯度
grads = (
grad_means3D,
grad_means2D,
grad_sh,
grad_colors_precomp,
grad_opacities,
grad_scales,
grad_rotations,
grad_cov3Ds_precomp,
None,
)
return grads
class GaussianRasterizationSettings(NamedTuple):
image_height: int
image_width: int
tanfovx : float
tanfovy : float
bg : torch.Tensor
scale_modifier : float
viewmatrix : torch.Tensor
projmatrix : torch.Tensor
sh_degree : int
campos : torch.Tensor
prefiltered : bool
debug : bool
高斯光栅化PyTorch模块
#用于高斯光栅化（Gaussian Rasterization）的PyTorch模块
class GaussianRasterizer(nn.Module): #定义了一个继承自nn.Module的类，表示高斯光栅化器。
#初始化方法，接受一个raster_settings参数，该参数包含了光栅化的设置（例如图像大小、视场、背景颜色等）。
def __init__(self, raster_settings):
super().__init__()
self.raster_settings = raster_settings
# 标记可见点的方法。接受3D点的位置作为输入，并使用C++/CUDA代码执行视锥体剔除，返回一个布尔张量，表示每个点是否可见。
def markVisible(self, positions):
# Mark visible points (based on frustum culling for camera) with a boolean
with torch.no_grad():
raster_settings = self.raster_settings
visible = _C.mark_visible(
positions,
raster_settings.viewmatrix,
raster_settings.projmatrix)
return visible
# 前向传播方法，用于进行高斯光栅化操作。接受一系列输入参数，包括3D坐标、2D坐标、透明度、SH特征或预计算的颜色、缩放、旋转或预计算的3D协方差等。
def forward(self, means3D, means2D, opacities, shs = None, colors_precomp = None, scales = None, rotations = None, cov3D_precomp = None):
raster_settings = self.raster_settings
# 检查SH特征和预计算的颜色是否同时提供，要求只提供其中一种。
if (shs is None and colors_precomp is None) or (shs is not None and colors_precomp is not None):
raise Exception('Please provide excatly one of either SHs or precomputed colors!')
# 检查缩放/旋转对或预计算的3D协方差是否同时提供，要求只提供其中一种。
if ((scales is None or rotations is None) and cov3D_precomp is None) or ((scales is not None or rotations is not None) and cov3D_precomp is not None):
raise Exception('Please provide exactly one of either scale/rotation pair or precomputed 3D covariance!')
# 如果某个输入参数为None，则将其初始化为空张量。
if shs is None:
shs = torch.Tensor([])
if colors_precomp is None:
colors_precomp = torch.Tensor([])
if scales is None:
scales = torch.Tensor([])
if rotations is None:
rotations = torch.Tensor([])
if cov3D_precomp is None:
cov3D_precomp = torch.Tensor([])
# 调用C++/CUDA光栅化例程rasterize_gaussians，传递相应的输入参数和光栅化设置。
# Invoke C++/CUDA rasterization routine
return rasterize_gaussians(
means3D,
means2D,
shs,
colors_precomp,
opacities,
scales,
rotations,
cov3D_precomp,
raster_settings,
)
高斯光栅化函数
这个函数调用了一个自定义的PyTorch Autograd Function _RasterizeGaussians.apply，并传递了一系列参数进行高斯光栅化。
def rasterize_gaussians(
means3D,
means2D,
sh,
colors_precomp,
opacities,
scales,
rotations,
cov3Ds_precomp,
raster_settings,
):
return _RasterizeGaussians.apply(
means3D, #高斯分布的三维坐标。
means2D, #高斯分布的二维坐标（屏幕空间坐标）。
sh, #SH（球谐函数）特征。
colors_precomp, #预计算的颜色。
opacities, #透明度
scales, #缩放因子
rotations, #旋转
cov3Ds_precomp, #预计算的三维协方差矩阵。
raster_settings, #高斯光栅化的设置。
)
rasterize_gaussians调用C++
代码路径：submodules\diff-gaussian-rasterization\ext.cpp
#include <torch/extension.h>
#include "rasterize_points.h"
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("rasterize_gaussians", &RasterizeGaussiansCUDA);
m.def("rasterize_gaussians_backward", &RasterizeGaussiansBackwardCUDA);
m.def("mark_visible", &markVisible);
}
以上3个C++方法的定义，在submodules\diff-gaussian-rasterization\rasterize_points.h文件中；
方法的具体实现在submodules\diff-gaussian-rasterization\cuda_rasterizer\rasterizer_impl.cu文件中。
涉及的前向渲染、后向渲染相关的部分，主要集中的submodules\diff-gaussian-rasterization\cuda_rasterizer，博客中有单独讲解，请阅读相关博客。
计算损失和反向传播
函数l1_loss、ssim的具体实现在文件utils\loss_utils.py。
报告或评估
训练报告
train.py中的训练报告：
def training_report(tb_writer, iteration, Ll1, loss, l1_loss, elapsed, testing_iterations, scene : Scene, renderFunc, renderArgs):
if tb_writer: #将 L1 loss、总体 loss 和迭代时间写入 TensorBoard。
tb_writer.add_scalar('train_loss_patches/l1_loss', Ll1.item(), iteration)
tb_writer.add_scalar('train_loss_patches/total_loss', loss.item(), iteration)
tb_writer.add_scalar('iter_time', elapsed, iteration)
# 在指定的测试迭代次数，进行渲染并计算 L1 loss 和 PSNR。
# Report test and samples of training set
if iteration in testing_iterations:
torch.cuda.empty_cache()
validation_configs = ({'name': 'test', 'cameras' : scene.getTestCameras()},
{'name': 'train', 'cameras' : [scene.getTrainCameras()[idx % len(scene.getTrainCameras())] for idx in range(5, 30, 5)]})
for config in validation_configs:
if config['cameras'] and len(config['cameras']) > 0:
l1_test = 0.0
psnr_test = 0.0
for idx, viewpoint in enumerate(config['cameras']):
# 获取渲染结果和真实图像
image = torch.clamp(renderFunc(viewpoint, scene.gaussians, *renderArgs)["render"], 0.0, 1.0)
gt_image = torch.clamp(viewpoint.original_image.to("cuda"), 0.0, 1.0)
if tb_writer and (idx < 5): # 在 TensorBoard 中记录渲染结果和真实图像
tb_writer.add_images(config['name'] + "_view_{}/render".format(viewpoint.image_name), image[None], global_step=iteration)
if iteration == testing_iterations[0]:
tb_writer.add_images(config['name'] + "_view_{}/ground_truth".format(viewpoint.image_name), gt_image[None], global_step=iteration)
# 计算 L1 loss 和 PSNR
l1_test += l1_loss(image, gt_image).mean().double()
psnr_test += psnr(image, gt_image).mean().double()
# 计算平均 L1 loss 和 PSNR
psnr_test /= len(config['cameras'])
l1_test /= len(config['cameras'])
# 在控制台打印评估结果
print("\n[ITER {}] Evaluating {}: L1 {} PSNR {}".format(iteration, config['name'], l1_test, psnr_test))
# 在 TensorBoard 中记录评估结果
if tb_writer:
tb_writer.add_scalar(config['name'] + '/loss_viewpoint - l1_loss', l1_test, iteration)
tb_writer.add_scalar(config['name'] + '/loss_viewpoint - psnr', psnr_test, iteration)
# 在 TensorBoard 中记录场景的不透明度直方图和总点数。
if tb_writer:
tb_writer.add_histogram("scene/opacity_histogram", scene.gaussians.get_opacity, iteration)
tb_writer.add_scalar('total_points', scene.gaussians.get_xyz.shape[0], iteration)
torch.cuda.empty_cache()#使用 torch.cuda.empty_cache() 清理 GPU 内存。
评估操作中的渲染设置
full_eval.py中有调用render.py
render.py主函数
if __name__ == "__main__":
# Set up command line argument parser
parser = ArgumentParser(description="Testing script parameters")
model = ModelParams(parser, sentinel=True)
pipeline = PipelineParams(parser)
parser.add_argument("--iteration", default=-1, type=int)
parser.add_argument("--skip_train", action="store_true")
parser.add_argument("--skip_test", action="store_true")
parser.add_argument("--quiet", action="store_true")
args = get_combined_args(parser)
print("Rendering " + args.model_path)
# Initialize system state (RNG)
safe_state(args.quiet)
render_sets(model.extract(args), args.iteration, pipeline.extract(args), args.skip_train, args.skip_test)
渲染设置集
def render_sets(dataset : ModelParams, iteration : int, pipeline : PipelineParams, skip_train : bool, skip_test : bool):
with torch.no_grad(): # 禁用梯度计算，因为在渲染过程中不需要梯度信息
gaussians = GaussianModel(dataset.sh_degree) # 创建一个 GaussianModel 对象，用于处理高斯模型
scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False) # 创建一个 Scene 对象，用于处理场景的渲染
bg_color = [1,1,1] if dataset.white_background else [0, 0, 0] # 根据数据集的背景设置，定义背景颜色
background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") # 将背景颜色转换为 PyTorch 张量，同时将其移到 GPU 上
if not skip_train: # 如果不跳过训练数据集的渲染
render_set(dataset.model_path, "train", scene.loaded_iter, scene.getTrainCameras(), gaussians, pipeline, background) # 调用 render_set 函数渲染训练数据集
if not skip_test: # 如果不跳过测试数据集的渲染
render_set(dataset.model_path, "test", scene.loaded_iter, scene.getTestCameras(), gaussians, pipeline, background) # 调用 render_set 函数渲染测试数据集
渲染设置
def render_set(model_path, name, iteration, views, gaussians, pipeline, background):
# 构建渲染结果和ground truth保存路径
render_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders")
gts_path = os.path.join(model_path, name, "ours_{}".format(iteration), "gt")
# 确保渲染结果和ground truth保存路径存在
makedirs(render_path, exist_ok=True)
makedirs(gts_path, exist_ok=True)
# 遍历所有视图进行渲染
for idx, view in enumerate(tqdm(views, desc="Rendering progress")):
# 调用 render 函数执行渲染，获取渲染结果
rendering = render(view, gaussians, pipeline, background)["render"] #这里执行的就是上面解析过的render的代码了~
# 获取视图的ground truth
gt = view.original_image[0:3, :, :]
# 保存渲染结果和ground truth为图像文件
torchvision.utils.save_image(rendering, os.path.join(render_path, '{0:05d}'.format(idx) + ".png"))
torchvision.utils.save_image(gt, os.path.join(gts_path, '{0:05d}'.format(idx) + ".png"))
参考资料：
https://blog.csdn.net/gwplovekimi/article/details/135500438
https://blog.csdn.net/qq_28087491/article/details/135629371
https://zhuanlan.zhihu.com/p/680669616
https://blog.csdn.net/qq_50791664/article/details/135932903
https://segmentfault.com/a/1190000045168206?decode__1660=eqmxn7qmqCqx9lDlxGrt7GOAIZ4Gw9ypD&u_atoken=7669f407ed1ee0eda7deec2089782e60&u_asig=2760821f17304309855405432e7ac1
原文地址：https://blog.csdn.net/xiner0114/article/details/143283944
免责声明：本站文章内容转载自网络资源，如侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！