Tip: 如果你在进行深度学习、自动驾驶、模型推理、微调或AI绘画出图等任务，并且需要GPU资源，可以考虑使用UCloud云计算旗下的Compshare的GPU算力云平台。他们提供高性价比的4090 GPU，按时收费每卡2.6元，月卡只需要1.7元每小时，并附带200G的免费磁盘空间。通过链接注册并联系客服，可以获得20元代金券（相当于6-7H的免费GPU资源）。欢迎大家体验一下~
0. 简介
我们所知道的是高斯溅射需要预训练，然后才能完成比较好的建图，这里提出的要求就是如何预训练出比较好的参数，然后再将关键帧对应的位置作为输入，并完成3D环境的重建。这里是一个使用VINS作为输入的整体流程。
1. 高斯溅射可视化
我们这里主要介绍的是最原始的高斯溅射模块。现在高斯溅射有一个比较常用的可视化软件
git clone https://github.com/graphdeco-inria/gaussian-splatting --recursive
# Dependencies
sudo apt install -y libglew-dev libassimp-dev libboost-all-dev libgtk-3-dev libopencv-dev libglfw3-dev libavdevice-dev libavcodec-dev libeigen3-dev libxxf86vm-dev libembree-dev
# Project setup
cd SIBR_viewers
cmake -Bbuild . -DCMAKE_BUILD_TYPE=Release # add -G Ninja to build faster
cmake --build build -j24 --target install
在安装后，我们可以使用SIBR_gaussianViewer_app，然后查看模型的效果
./SIBR_viewers/install/bin/SIBR_gaussianViewer_app -m /home/linhai/xxx/gaussian-splatting/output/ba17fb80-1
2. 高斯溅射代码
这里搬运了gwpscut的高斯溅射文章，首先就是看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)
其定义如下：
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"))
这段代码定义了一个函数 safe_state(silent)，该函数的作用是在执行期间重定向标准输出（sys.stdout）到一个新的类 F 的实例。这个类 F 在写入时会检查是否需要在每行结尾处添加时间戳，以及是否需要替换换行符。
具体来说，函数的实现步骤如下：
- 将原始的标准输出保存在
old_f变量中。 - 定义一个名为
F的新类，该类具有以下方法：__init__(self, silent)：初始化方法，接受一个参数silent。write(self, x)：写入方法，检查silent属性，如果不是静默模式，则在每行结尾添加当前时间戳，并将文本写入原始标准输出。flush(self)：刷新方法，将原始标准输出的缓冲区刷新。
- 创建
F类的实例并将其赋值给sys.stdout，从而重定向标准输出到新的类实例。 - 设置随机种子以确保结果的可重复性。
- 最后，将 PyTorch 的随机种子设置为 0，并将当前 CUDA 设备设置为 “cuda:0”（如果可用的话）。
这段代码的目的是为了在执行过程中控制标准输出的行为，添加时间戳并在需要时禁止输出，以便在某些场景下更方便地进行调试和记录。
（应该就是输出一些系统的状态的）
然后就是启动GUI以及运行训练的代码
# Start GUI server, configure and run training
network_gui.init(args.ip, args.port) #这行代码初始化一个 GUI 服务器，使用 args.ip 和 args.port 作为参数。这可能是一个用于监视和控制训练过程的图形用户界面的一部分。
torch.autograd.set_detect_anomaly(args.detect_anomaly) #这行代码设置 PyTorch 是否要检测梯度计算中的异常。
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)
# 输入的参数包括：模型的参数（数据集的位置）、优化器的参数、其他pipeline的参数，测试迭代次数、保存迭代次数 、检查点迭代次数 、开始检查点 、调试起点
接下来就是training这个函数到底做了什么了（具体的解析过程见注释，由于代码不少，没办法一句一句顺思路来解析）
def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoint_iterations, checkpoint, debug_from):
first_iter = 0 #初始化迭代次数。
tb_writer = prepare_output_and_logger(dataset) #设置 TensorBoard 写入器和日志记录器。
gaussians = GaussianModel(dataset.sh_degree) #（重点看，需要转跳）创建一个 GaussianModel 类的实例，输入一系列参数，其参数取自数据集。
scene = Scene(dataset, gaussians) #（这个类的主要目的是处理场景的初始化、保存和获取相机信息等任务，）创建一个 Scene 类的实例，使用数据集和之前创建的 GaussianModel 实例作为参数。
gaussians.training_setup(opt) #设置 GaussianModel 的训练参数。
if checkpoint: #如果有提供检查点路径。
(model_params, first_iter) = torch.load(checkpoint)#通过 torch.load(checkpoint) 加载检查点的模型参数和起始迭代次数。
gaussians.restore(model_params, opt)#通过 gaussians.restore 恢复模型的状态。
bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0] #设置背景颜色，根据数据集是否有白色背景来选择。
background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") #将背景颜色转化为 PyTorch Tensor，并移到 GPU 上。
# 创建两个 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
progress_bar = tqdm(range(first_iter, opt.iterations), desc="Training progress") #创建一个 tqdm 进度条，用于显示训练进度。
first_iter += 1
# 接下来开始循环迭代
for iteration in range(first_iter, opt.iterations + 1): #主要的训练循环开始。
if network_gui.conn == None: #检查 GUI 是否连接，如果连接则接收 GUI 发送的消息。
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
if iteration % 1000 == 0:
gaussians.oneupSHdegree() #每 1000 次迭代，增加球谐函数的阶数。
# 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
gt_image = viewpoint_cam.original_image.cuda()
Ll1 = l1_loss(image, gt_image)
loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim(image, gt_image)) #计算渲染的图像与真实图像之间的loss
loss.backward() #更新损失。loss反向传播
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
gaussians.max_radii2D[visibility_filter] = torch.max(gaussians.max_radii2D[visibility_filter], radii[visibility_filter])
gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter)
if iteration > opt.densify_from_iter and iteration % opt.densification_interval == 0:
size_threshold = 20 if iteration > opt.opacity_reset_interval else None
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):
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")
接下来要看的就是高斯模型的定义
gaussians = GaussianModel(dataset.sh_degree) #创建一个 GaussianModel 类的实例，输入一系列参数，其参数取自数据集。
解读转跳下一节。
而损失函数的定义，如下
# Loss
gt_image = viewpoint_cam.original_image.cuda()
Ll1 = l1_loss(image, gt_image)
loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim(image, gt_image)) #计算渲染的图像与真实图像之间的loss
loss.backward() #更新损失。loss反向传播
也正是参数优化的损失函数（L1 与 D-SSIM 项的组合）
总结训练过程每个迭代主要执行以下的操作：
- 每 1000 次迭代，增加球谐系数的阶数。
- 随机选择一个相机视角。
- 渲染图像，获取视点空间点、能见度过滤器和半径等信息。
- 计算损失（L1 损失和 DSSIM 损失的加权和），进行反向传播。
- 通过无梯度的上下文进行后续操作：
- 根据迭代次数进行点云密度操作（densification）：
- 更新最大半径信息。
- 根据条件进行点云密度增加和修剪。
- 进行优化器的参数更新。
在整个训练过程中，这些步骤循环执行，逐渐优化模型参数，进行损失计算和反向传播，同时根据条件进行点云密度操作和保存检查点，以逐步提升模型性能。
2.1 高斯模型的定义
高斯模型的初始化，初始化过程中加载或定义了各种相关的属性使用的球谐阶数、最大球谐阶数、各种张量（_xyz、_features_dc等）、优化器和其他参数。
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 方法设置各种激活和变换函数
调用 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 #用于归一化旋转矩阵。
高斯模型初始化后在下面代码中进行加载（从已加载的模型中加载还是从点云数据中加载）
# 在train.py中
scene = Scene(dataset, gaussians) #（这个类的主要目的是处理场景的初始化、保存和获取相机信息等任务，）创建一个 Scene 类的实例，使用数据集和之前创建的 GaussianModel 实例作为参数。
#对应的在Scene类中
# 加载或创建高斯模型
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)
2.2 自适应密集化处理
在train.py中，在一定的迭代次数内进行密集化处理（这也是3DGS的一大贡献点，自适应控制密集化）。见下面代码
# Densification（在一定的迭代次数内进行密集化处理。）
if iteration < opt.densify_until_iter: #在达到指定的迭代次数之前执行以下操作。
# Keep track of max radii in image-space for pruning
gaussians.max_radii2D[visibility_filter] = torch.max(gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) #将每个像素位置上的最大半径记录在 max_radii2D 中。这是为了密集化时进行修剪（pruning）操作时的参考。
gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) #将与密集化相关的统计信息添加到 gaussians 模型中，包括视图空间点和可见性过滤器。
if iteration > opt.densify_from_iter and iteration % opt.densification_interval == 0: #在指定的迭代次数之后，每隔一定的迭代间隔进行以下密集化操作。
size_threshold = 20 if iteration > opt.opacity_reset_interval else None #根据当前迭代次数设置密集化的阈值。如果当前迭代次数大于 opt.opacity_reset_interval，则设置 size_threshold 为 20，否则为 None。
gaussians.densify_and_prune(opt.densify_grad_threshold, 0.005, scene.cameras_extent, size_threshold) #执行密集化和修剪操作，其中包括梯度阈值、密集化阈值、相机范围和之前计算的 size_threshold。
if iteration % opt.opacity_reset_interval == 0 or (dataset.white_background and iteration == opt.densify_from_iter): #在每隔一定迭代次数或在白色背景数据集上的指定迭代次数时，执行以下操作。
gaussians.reset_opacity() #重置模型中的某些参数，涉及到透明度的操作，具体实现可以在 reset_opacity 方法中找到。
而进入下面函数gaussians.densify_and_prune（）解读如下：
# 执行密集化和修剪操作
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函数
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)
其中，最后一句代码实现的是将新的密集化点的相关特征保存在一个字典中。
# 将新的密集化点的相关特征保存在一个字典中。
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}
optimizable_tensors = self.cat_tensors_to_optimizer(d) #将字典中的张量连接（concatenate）成可优化的张量。这个方法的具体实现可能是将字典中的每个张量进行堆叠，以便于在优化器中进行处理。
# 更新模型中原始点集的相关特征，使用新的密集化后的特征。
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")
对于densify_and_split函数
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)
2.3 渲染（render）及光栅化
train.py在求loss前，是需要进行图像的渲染并且计算loss的，见下代码：
# 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
gt_image = viewpoint_cam.original_image.cuda()
Ll1 = l1_loss(image, gt_image)
loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim(image, gt_image)) #计算渲染的图像与真实图像之间的loss
loss.backward() #更新损失。loss反向传播
接下来就看看渲染的过程。
由于3DGS的渲染过程是利用了光栅化（rasterization）而光栅化的过程需要在GPU上运行。
# 这段代码是一个用于渲染场景的函数，主要是通过将高斯分布的点投影到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}
3DGS看了这么久，知道SH系数是跟颜色相挂钩的，至于如何相挂钩好像原文是没有细说的，但是代码实现就很直观了
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之间。
在这段代码中其中就调用了submodules的diff_gaussian_rasterization
from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
这部分的源码用cuda写的（这也是本人第一次真正深入看cuda源码~）
GaussianRasterizationSettings与GaussianRasterizer对应的代码在submodules/diff-gaussian-rasterization/diff_gaussian_rasterization/__init__.py
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
class GaussianRasterizer(nn.Module):
def __init__(self, raster_settings):
super().__init__()
self.raster_settings = raster_settings
而上面进行的都是光栅的初始化，真正继续渲染计算应该是下面代码（输入具体的数值）
# 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)
接下来学习一个光栅化的源码。主要的运行及计算的工程应该就是forward（采用前向渲染）
# 用于高斯光栅化（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,
)
而最后的rasterize_gaussians的定义如下
def rasterize_gaussians( #这个函数调用了一个自定义的PyTorch Autograd Function _RasterizeGaussians.apply，并传递了一系列参数进行高斯光栅化。
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, #高斯光栅化的设置。
)
它又return一个_RasterizeGaussians
# 这是一个自定义的 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
而cuda的具体实现则在下面代码中
先看前向的渲染，进入forward.h文件可以发现，其实就是两个函数，预处理与渲染。
先看预处理，该函数使用了CUDA并行计算，通过调用名为 preprocessCUDA 的 CUDA 核函数来执行高斯光栅化的前处理。CUDA 核函数的执行由函数参数确定。在 CUDA 核函数中，每个线程块由多个线程组成，负责处理其中的一部分数据，从而加速高斯光栅化的计算。
125
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
