3D Gaussian Splatting代码中的train和render两个文件代码解读

现在来聊一聊训练和渲染是如何进行的

training

train.py
line 31
def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoint_iterations, checkpoint, debug_from):# 初始化第一次迭代的索引为0first_iter = 0# 准备输出和日志记录器tb_writer = prepare_output_and_logger(dataset)# 初始化高斯模型，参数为数据集的球谐函数（SH）级别gaussians = GaussianModel(dataset.sh_degree)# 创建场景对象，包含数据集和高斯模型scene = Scene(dataset, gaussians)# 设置高斯模型的训练配置gaussians.training_setup(opt)# 加载检查点（如果有），恢复模型参数和设置起始迭代次数if checkpoint:(model_params, first_iter) = torch.load(checkpoint)gaussians.restore(model_params, opt)# 设置背景颜色，如果数据集背景为白色，则设置为白色（[1, 1, 1]），否则为黑色（[0, 0, 0]）bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]# 将背景颜色转换为CUDA张量，以便在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# 用于记录指数移动平均损失的变量，初始值为0.0ema_loss_for_log = 0.0# 创建进度条，用于显示训练进度，从起始迭代数到总迭代数progress_bar = tqdm(range(first_iter, opt.iterations), desc="Training progress")# 增加起始迭代数，以便从下一次迭代开始first_iter += 1for iteration in range(first_iter, opt.iterations + 1):# 尝试连接网络GUI，如果当前没有连接if network_gui.conn == None:network_gui.try_connect()# 如果已经连接网络GUI，处理接收和发送数据while network_gui.conn != None:try:# 初始化网络图像字节为Nonenet_image_bytes = None# 从网络GUI接收数据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())# 发送渲染结果到网络GUI，并附带数据集的源路径network_gui.send(net_image_bytes, dataset.source_path)# 如果需要进行训练，并且当前迭代次数小于总迭代次数，或不需要保持连接，则退出循环if do_training and ((iteration < int(opt.iterations)) or not keep_alive):breakexcept Exception as e:# 如果出现异常，断开网络连接network_gui.conn = None# 记录当前迭代的开始时间，用于计算每次迭代的持续时间iter_start.record()# 更新学习率gaussians.update_learning_rate(iteration)# 每1000次迭代增加一次SH级别，直到达到最大度if iteration % 1000 == 0:gaussians.oneupSHdegree()# 随机选择一个相机视角if not viewpoint_stack:viewpoint_stack = scene.getTrainCameras().copy()# 从相机视角堆栈中随机弹出一个相机视角viewpoint_cam = viewpoint_stack.pop(randint(0, len(viewpoint_stack) - 1))# 渲染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"]# 计算损失gt_image = viewpoint_cam.original_image.cuda()  # 获取地面真实图像Ll1 = l1_loss(image, gt_image)  # 计算L1损失# 计算总损失，结合L1损失和结构相似性损失（SSIM）loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim(image, gt_image))loss.backward()  # 反向传播计算梯度# 记录当前迭代的结束时间，用于计算每次迭代的持续时间iter_end.record()# 在不需要计算梯度的上下文中进行操作with torch.no_grad():# 更新进度条和日志ema_loss_for_log = 0.4 * loss.item() + 0.6 * ema_loss_for_log  # 更新指数移动平均损失if iteration % 10 == 0:# 每10次迭代更新一次进度条progress_bar.set_postfix({"Loss": f"{ema_loss_for_log:.{7}f}"})progress_bar.update(10)if iteration == opt.iterations:progress_bar.close()# 记录训练报告并保存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)# 密集化操作if iteration < opt.densify_until_iter:# 跟踪图像空间中的最大半径，用于修剪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 Nonegaussians.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()# 优化器步骤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")

render

现在是渲染的这个文件进行方式,首先是主文件里单张图片的渲染和整个数据集的渲染方法：

render.py
line 24
# 渲染一组视角并保存渲染结果和对应的真实图像
def render_set(model_path, name, iteration, views, gaussians, pipeline, background):# 定义渲染结果和真实图像的保存路径render_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders")gts_path = os.path.join(model_path, name, "ours_{}".format(iteration), "gt")# 创建保存路径，如果路径不存在makedirs(render_path, exist_ok=True)makedirs(gts_path, exist_ok=True)# 遍历每个视角进行渲染for idx, view in enumerate(tqdm(views, desc="Rendering progress")):# 渲染图像rendering = render(view, gaussians, pipeline, background)["render"]# 获取对应的真实图像gt = view.original_image[0:3, :, :]# 保存渲染结果和真实图像到指定路径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"))# 渲染训练集和测试集的图像，并保存结果
def render_sets(dataset: ModelParams, iteration: int, pipeline: PipelineParams, skip_train: bool, skip_test: bool):with torch.no_grad():# 初始化高斯模型和场景gaussians = GaussianModel(dataset.sh_degree)scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False)# 设置背景颜色bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")# 如果不跳过训练集渲染，则渲染训练集的图像if not skip_train:render_set(dataset.model_path, "train", scene.loaded_iter, scene.getTrainCameras(), gaussians, pipeline, background)# 如果不跳过测试集渲染，则渲染测试集的图像if not skip_test:render_set(dataset.model_path, "test", scene.loaded_iter, scene.getTestCameras(), gaussians, pipeline, background)

但是这两个方法都是外层函数，并没有展示渲染如何进行参数传递和具体操作，在以下代码中才是最关键的内容：

gaussian_renderer\__init__.py
line 18
def render(viewpoint_camera, pc: GaussianModel, pipe, bg_color: torch.Tensor, scaling_modifier=1.0, override_color=None):"""渲染场景。参数:viewpoint_camera - 摄像机视角pc - 高斯模型pipe - 管道参数bg_color - 背景颜色张量，必须在GPU上scaling_modifier - 缩放修饰符，默认为1.0override_color - 覆盖颜色，默认为None"""# 创建一个全零张量，用于使PyTorch返回2D（屏幕空间）均值的梯度screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0try:screenspace_points.retain_grad()  # 保留梯度信息except:pass# 设置光栅化配置tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)  # 计算视角的X轴正切tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)  # 计算视角的Y轴正切raster_settings = GaussianRasterizationSettings(image_height=int(viewpoint_camera.image_height),  # 图像高度image_width=int(viewpoint_camera.image_width),  # 图像宽度tanfovx=tanfovx,  # 视角X轴正切tanfovy=tanfovy,  # 视角Y轴正切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  # 获取3D均值means2D = screenspace_points  # 获取2D均值opacity = pc.get_opacity  # 获取不透明度# 如果提供了预计算的3D协方差，则使用它。如果没有，则从光栅化器的缩放/旋转中计算。scales = Nonerotations = Nonecov3D_precomp = Noneif pipe.compute_cov3D_python:cov3D_precomp = pc.get_covariance(scaling_modifier)  # 计算3D协方差else:scales = pc.get_scaling  # 获取缩放rotations = pc.get_rotation  # 获取旋转# 如果提供了预计算的颜色，则使用它们。否则，如果需要在Python中预计算SH到颜色的转换，则进行转换。shs = Nonecolors_precomp = Noneif 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)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)colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)  # 计算颜色else:shs = pc.get_features  # 获取球谐函数特征else:colors_precomp = override_color  # 覆盖颜色# 将可见的高斯体光栅化为图像，并获取它们在屏幕上的半径。rendered_image, radii = rasterizer(means3D=means3D,means2D=means2D,shs=shs,colors_precomp=colors_precomp,opacities=opacity,scales=scales,rotations=rotations,cov3D_precomp=cov3D_precomp)# 那些被视锥剔除或半径为0的高斯体是不可见的。# 它们将被排除在用于分裂标准的值更新之外。return {"render": rendered_image,  # 渲染图像"viewspace_points": screenspace_points,  # 视图空间点"visibility_filter": radii > 0,  # 可见性过滤器"radii": radii  # 半径}