基于昇腾平台的 3D Gaussian Splatting 训推优化实践：环境搭建、NPU 算子编译与融合算子加速

基于昇腾平台的 3D Gaussian Splatting 训推优化实践环境搭建、NPU 算子编译与融合算子加速【免费下载链接】cann-recipes-embodied-ai本项目针对具身智能业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-ai3D Gaussian Splatting3DGS是 2023 年提出的可微渲染算法以 3D 高斯椭球为核心图元通过可微光栅化完成三维场景重建与渲染打破了传统建模速度与保真度不可兼得的瓶颈。本文以 3DGS 昇腾适配样例 为核心完整讲解在 CANN 昇腾 Atlas A2 单卡环境下的环境搭建、自定义 NPU 算子编译安装、数据集准备与训练/推理全流程并结合仓库内的四篇优化专题文档与算子源码深入剖析 Alpha-blending、视锥剔除、负载均衡、Precise Intersection 四类融合算子的实现原理与性能收益帮助读者掌握在 SIMD 架构上移植并优化 3DGS 训推链路的完整实战方案。项目概览四大 NPU 优化点本项目提供 3DGS 的昇腾适配版本整体基于 gsplat 开源框架改造训练与推理引擎位于 rasterization 目录。相比 GPU 实现针对昇腾 SIMD 计算架构本项目主要实现了以下优化点详细原理见 NPU 3DGS 训推优化实践Alpha-blending 优化算法及融合算子实现优化通过指令流水并行、Double Buffer 与两级剪枝机制加速逐像素颜色累积视锥剔除融合算子优化将投影前的多类剔除逻辑融合为单个GaussianFilter算子避免耗时的取 index tensor move 操作Gaussian 负载均衡优化针对像素块间高斯球数量分布不均的问题采用贪心算法分配负载消除单核拖尾Precise Intersection 融合算子优化基于 FlashGS 算法以更精确的椭圆相切矩形与相交判断替代传统正切四边形映射减少无效高斯球的排序与渲染。这四个优化点均以自定义 Ascend C 算子形式落地通过meta_gauss_render算子库封装为 Python 可调用接口最终由 rasterizer.py 串联成完整的 3DGS 渲染流水。执行环境与 CANN 环境准备本样例支持昇腾 Atlas A2 环境的单卡训练和推理。运行前需按顺序完成 CANN 软件、PyTorch 插件与自定义算子的安装。CANN 开发套件与二进制算子包样例依赖 CANN 开发套件包cann-toolkit与 CANN 二进制算子包cann-kernels当前使用的 CANN 软件版本为CANN 8.2.RC1。请从 CANN 社区下载中心获取对应版本选择与操作系统架构匹配的Ascend-cann-toolkit_${version}_linux-${arch}.run与Ascend-cann-kernels-${chip_type}_${version}_linux-${arch}.run两个软件包并参考 CANN 安装文档以本地安装方式完成安装。本样例依赖的torch 与 torch_npu 版本分别为 2.1.0 和 2.1.0.post12请从昇腾官方 Ascend Extension for PyTorch 插件下载页获取并安装对应版本。随后创建并激活 Python 3.8 的 conda 环境conda create -n 3dgs python3.8 conda activate 3dgs自定义 NPU 算子编译安装本项目依赖的自定义算子需要以 whl 包形式编译安装进当前 conda 环境整体工程位于 ops/ascendc。推荐使用gcc 10.2版本编译。如需安装编译期依赖执行以下指令pip install numpy1.23 decorator sympy scipy attrs cloudpickle psutil synr0.5.0 tornado cmake pyyaml expecttest protobuf克隆原始仓git clone https://gitcode.com/cann/cann-recipes-embodied-ai.git编译自定义 NPU 算子cd ops/ascendc bash build.sh --python3.8参数--python指定编译使用的 Python 版本支持 3.8 及以上版本。编译成功后会在ops/ascendc目录下生成build、meta_gauss_render.egg-info、dist文件夹生成的 whl 包位于dist目录下。安装 NPU 算子在ops/ascendc路径下执行pip install dist/*.whl --force-reinstall该算子包对应 meta_gauss_render 模块其 csrc 目录下的CalcRender.cpp、GaussianFilter.cpp、GaussianSort.cpp、FlashGaussianBuildMask.cpp、ProjectionThreeDimsGaussianForward.cpp、QuatScalesToCovars.cpp、SphericalHarmonics*.cpp等文件正是四大优化算子与投影、排序等基础算子的昇腾实现ops 目录则提供对应的 Python 封装接口。Python 外部依赖与数据集准备进入样例目录并安装依赖清单中的第三方库cd ../../3d_vision/gaussian_splatting pip install -r requirements.txt数据集通过 download_dataset.py 一键下载Mip-NeRF 360 数据集的 360_v2 子集在3d_vision/gaussian_splatting路径下执行python datasets/download_dataset.py数据集会被下载并解压到3d_vision/gaussian_splatting/data/360_v2路径下例如garden、bonsai等场景目录每个场景内含 COLMAP 格式的位姿、图像与稀疏点云作为 3DGS 初始化的 SfM 种子点由 datasets/colmap.py 解析。快速启动训练与推理启动前需按 CANN 安装教程配置环境变量source /usr/local/Ascend/ascend-toolkit/set_env.sh运行训练脚本ASCEND_RT_VISIBLE_DEVICES1 python train.pyASCEND_RT_VISIBLE_DEVICES1指定使用 NPU 设备 1 卡不指定则默认启动 0 卡。模型训推的配置文件为 rasterization/config.py--data_dir指定训练数据集路径默认为data/360_v2/garden--result_dir指定训练结果输出路径默认为results/garden。换用其他场景的示例python train.py --data_dir data/360_v2/bonsai --result_dir results/bonsai运行推理脚本推理与训练共用 train.py 入口传入--ckpt参数指定训练保存的模型如results/bonsai/ckpts/ckpt_29999_rank0.pt即跳过训练、对输入数据集进行渲染与评估python train.py --data_dir data/360_v2/bonsai --result_dir results/bonsai --ckpt path_to_ckpt.pt从 train.py 可以看到当cfg.ckpt非空时Runner 只执行runner.eval(stepstep)完成推理评估否则执行runner.train()进入训练循环。数据集与评估模型的场景类别需要保持一致推理与评估结果渲染图、PSNR/SSIM/LPIPS 指标、统计 json保存在result_dir目录下。训练所有场景数据集train_all.sh 按脚本中指定顺序执行数据集中所有场景的训练和推理bash train_all.sh配置文件详解训练与推理的全部超参数集中在 rasterization/config.py 的Config数据类中通过 train.py 中的tyro.cli(Config)解析为命令行参数并支持--steps_scaler对训练步数相关参数进行整体缩放。核心参数说明如下参数默认值说明ckptNone已训练模型.pt文件路径列表传入则跳过训练仅执行评估data_dirdata/360_v2/gardenMip-NeRF 360 数据集路径data_factor4数据集下采样倍数result_dirresults/garden结果保存目录含 ckpts / stats / renders / ply / tbtest_every8每隔 N 张图像取 1 张作为测试图像tile_size32像素块边长可选 32 或 64直接决定渲染算子的 tile 切分粒度batch_size1训练批大小学习率会被自动缩放max_steps30000总训练步数eval_steps[500, 7000, 30000]执行评估的步数集合save_steps[7000, 30000]保存.pt检查点的步数集合save_ply/ply_stepsFalse/[7000, 30000]是否及何时导出 PLY 点云init_typesfm高斯初始化策略sfm使用 SfM 点云init_num_pts100000非 SfM 初始化时的高斯初始数量sh_degree3球谐函数最高阶数渲染时按步数渐进启用间隔sh_degree_interval1000init_opa/init_scale0.1/1.0高斯初始不透明度与尺度ssim_lambda0.2SSIM 损失权重总损失为L1 * (1 - λ) SSIM * λ见 runner.pynear_plane/far_plane0.01/1e10近/远裁剪面距离opacity_reg/scale_reg0.0不透明度 / 尺度正则权重depth_loss/depth_lambdaFalse/1e-2是否启用深度损失及权重在视差空间计算见 runner.pylpips_netalexLPIPS 骨干网络可选alex或vgg训练过程中每 800 步会在renders目录写入 GT 与渲染的拼接图达到eval_steps时自动在验证集上计算 PSNR、SSIM、LPIPS 并打印与落盘。目录结构与渲染调用链3d_vision/gaussian_splatting/ ├── datasets/ # 数据集下载及解析模块colmap.py、download_dataset.py、normalize.py ├── gsplat/ # gsplat 框架提供的工具函数和分布式训练模块utils.py、distributed.py ├── rasterization/ # 算法主要实现 │ ├── config.py # 模型训练、推理配置文件 │ ├── utils.py # 渲染和模型训练相关工具函数高斯创建、优化器、SSIM 等 │ ├── runner.py # 训练和推理引擎Runner 类 │ └── rasterizer.py # 3DGS 渲染流程实现内部调用并串联各个自定义算子 ├── README.md # 项目说明 ├── requirements.txt # 外部依赖 ├── train.py # 训练和推理启动脚本 └── train_all.sh # 所有场景的训推拉起脚本从源码看rasterizer.py 从meta_gauss_render导入spherical_harmonics、projection_three_dims_gaussian_fused、flash_gaussian_build_mask、gaussian_sort、calc_render、get_render_schedule六个算子_ascend_rasterization方法完整呈现了昇腾渲染流水rasterizer.py球谐颜色解码spherical_harmonics根据视线方向计算 RGB 颜色融合投影projection_three_dims_gaussian_fused一步完成 3D→2D 投影输出 means2d、深度、conics逆协方差、2D 协方差等tile 映射与排序flash_gaussian_build_mask计算高斯球与像素块的映射Precise Intersection与 tile 内深度排序所需信息负载调度get_render_schedule在 CPU 侧基于每 tile 高斯计数执行贪心负载均衡生成调度张量lb_sched排序gaussian_sort按 tile 内深度序重排高斯 ID光栅化calc_render逐相机调用完成 Alpha blending 渲染出颜色与深度图再经tile2image重排回图像空间。其中投影、排序与 Alpha blending 对应的自定义算子源码位于 ops/ascendc/kernels/op_kernel如projection_three_dims_gaussian_forward.cpp、gaussian_sort.cpp、calc_render_fwd_double_clip_gsids.cpp配套的逐算子精度测试位于 ops/ascendc/tests可供二次开发与验证参考。优化一Alpha-blending 融合算子与流水剪枝设计Alpha blending 是 3DGS 光栅化的核心步骤给定相机视角与投影平面将 3D 高斯球的颜色表达渲染到各像素上。计算公式为$$C \sum_{i \in N} c_i \alpha_i \prod^{i-1}_{j1}(1-\alpha_j)$$$$\alpha_i o_i \exp(-\frac{1}{2} \delta^T \Sigma \delta)$$整体流程为按 tile 切分 vector core每个核循环渲染分配到的 tile 上所有像素由于累乘依赖必须对高斯球按深度排序后依序渲染对每个高斯球并行计算其对 tile 所有像素的 α 值与颜色贡献再逐像素更新累积颜色与透明度。在 SIMD 架构下实现该算法存在三大难点流水设计复杂需在 vector 指令间频繁插入 PipeBarrier 保证数据正确性内存访问与计算开销难以互相掩盖tile 内渲染的高斯球在内存中不连续每个高斯球需对 tile 上所有像素计算冗余计算量大。指令流水并行根据 msprof op simulator 仿真结果对 64 个 float 计算 VMuls 需 15ns而 2048 个 float 的 VMuls 仅需 51ns指令解码与数据预取开销被摊薄。基于此项目将逐高斯球的串行流水改为每批处理 4 个高斯球后再插入一次 PipeBarrier先把单个高斯球的计算拆成两阶段——第一阶段并行计算 4 个高斯球的 α 值并暂存第二阶段依次计算颜色贡献并更新累积值。该设计显著减少 PipeBarrier 插入频率代价是 UB 空间占用上升需仔细复用中间缓冲如颜色渲染时复用前一高斯球 α 计算的存储空间。完整流水示意见 流水图。Double Buffer 优化为掩盖内存访问开销采用双缓冲提前读取高斯球属性处理当前批次高斯球的同时把下一批次高斯球属性预取到第二块 buffer 中处理时直接读 buffer避免内存访问延迟见 pingpong buffer 示意。两级剪枝设计渲染中大量高斯球对当前 tile 无贡献项目设计了两级剪枝Tile 级别剪枝当 tile 累积透明度达到阈值时提前结束该 tile 渲染。受 SIMD 架构限制无法逐像素剪枝因此每次处理完一组4 个高斯球后用CompareScalar将全部像素累积透明度与阈值比较生成掩码ReinterpretCast转为 64 位整数再与UINT64_MAX按位与判断是否全部达标同时记录每个 tile 终止渲染的高斯球 ID供反向传播从该位置续算梯度。该策略性能收益约 5%~10%效果有限。Sub Tile像素行级别剪枝通过协方差矩阵计算高斯球在投影平面的影响范围。实验表明每个高斯球实际只影响约 6×6 像素而预设 tile 为 32×32即只影响约 20% 的像素行。实现上在整个 tile 上计算 α 值确定覆盖的像素行范围用WholeReduceMax求每行 α 最大值CompareScalar与阈值比较生成掩码并解释为 64 位整数用ScalarCountLeadingZeros与ScalarGetCountOfValue1确定需要渲染的像素行区间注意小端存储下首个 1 对应最后一个需渲染的行并将每个高斯球实际影响的起始/结束行索引以 UInt8 存于 UB每 1024 个写回 GM 供反向使用见 像素行剪枝示意图。实验数据在约 13 万高斯球过滤后实际渲染 67917 个的场景训练 30000 步α 阈值取 0.004、不透明度阈值取 0.99优化方法PSNR前向 device 耗时(ms)反向 device 耗时(ms)无优化25.0084.80718.710流水优化25.0054.37917.565剪枝优化25.0105.10113.141流水 剪枝24.9744.49513.107正向耗时收益约 10%反向耗时收益约 30%。不同 α 阈值仅评估未重训的影响阈值提高对反向收益明显取 0.1 时反向 device 耗时仅为原来的 7.5%而 PSNR 只下降 3.3%25.082→24.274取 0.2 时反向耗时进一步降至 1.043ms但 PSNR 降至 23.309。优化二视锥剔除融合算子GaussianFilterCulling剔除是渲染前的关键优化步骤只保留当前相机能看到的高斯减少后续光栅化与 Alpha blending 计算量。投影预处理阶段的视锥体剔除逻辑包括无效高斯剔除确保高斯是合法的正定矩阵非退化、非扁平深度方向剔除剔除深度小于近裁剪面或大于远裁剪面的高斯水平/垂直方向剔除剔除投影后 2D 椭圆 x/y 方向半径为 0 的高斯屏幕外剔除仅保留投影椭圆与屏幕有重叠的高斯。实现上分两步一是实现GaussianFilter将视锥剔除融合为单个算子避免耗时的取 index tensormove 操作返回比特位掩码filter1 有效、0 无效二是投影预处理反向算子计算前做反视锥剔除还原原始 Tensor无效高斯梯度自然置 0。对应实现位于 op_host/gaussian_filter.cpp 与 op_kernel/gaussian_filter.cpp。实现难点有三int8 每个元素存 8 个高斯掩码状态分片必须是 8 的整数倍否则比特位跨元素映射混乱算子切分涉及核间数据依赖需引入全局同步但又要控制同步次数反向需把过滤数据分散还原到原始位置核心是为无规则filter构建GatherMask数据分散索引。tiling 设计高斯数量普遍在十万以上故不对 batchsize 与相机数分核而是对高斯分核。filter 为 int8头核需 8 元素对齐对齐后尾核可能算出负数需反推实际使用的核数再计算尾核。设核数为coreNum、高斯总数为N、头核处理perCoreN、尾核处理lastCoreN$$perCoreN \lceil \lceil N / coreNum \rceil / 8 \rceil \cdot 8$$$$usedCoreNum \lceil N / perCoreN \rceil$$$$lastCoreN N - usedCoreNum \cdot perCoreN$$示意见 tiling 分核图kernel 偏移计算受 UB 限制每个核一次迭代无法过滤出全部有效高斯需记录每核有效高斯数cntPerCore_i与每次迭代有效数cntPerCore_{i,j}则第 i 核第 j 次迭代的搬出偏移为$$offset \sum_{k0}^{i-1}{cntPerCore_k} \sum_{k0}^{j-1}{cntPerCore_{i,k}}$$计算流程偏移依赖前序核结果而各核并行故用SyncAll分两阶段。Phase1 执行loopN次循环搬入、计算、搬出 filter同时统计核内迭代计数数组cntPerLoop与总计数cntPerCore结束后把cntPerCore存 workspace全局同步迁出 GM 上的核间输出偏移coreOffset仅最后一个核汇总全局计数。Phase2 重新搬入 filter以 filter 为 mask 用GatherMask执行剔除结合coreOffset与cntPerLoop确定全局偏移用 3 个TQue保证 MTE2/Vector/MTE3 流水并行。反向 Scatter 实现为避免循环 标量计算的耗时用GatherMaskSortGather组合实现数据分散先构造等长元素索引Index分别按 filter 位为 1/0 收集得到headIndexTensor/tailIndexTensor拼接注意 32 字节对齐先拷到 GM 连续地址再拷回 UB成新索引newIndex再经Sort排序并Extract分离排序值与其原位置索引由于Sort只支持降序需构造等差数列重排得到升序索引最后用gather按索引收集还原数据。数据搬运方面用Transpose接口替代构造 index 的GatherMask搬运如(B,3,N)搬运为(B,N,3)避免大量标量运算构建 index。优化三Gaussian 负载均衡策略3DGS Alpha blending 按像素块对 vector core 均匀切分时各像素块投影的高斯球数量分布差异较大会导致个别核计算拖尾、整体渲染耗时增加。GPU 实现可依赖硬件 Warp 级调度而 SIMD 架构必须显式处理负载不均。项目采用贪心算法做负载均衡目标是让各 vector 核负载尽可能均衡遵循先重后轻、分给最轻原则先将像素块按高斯球数从大到小排序再把排序后的像素块依次分配给当前累积高斯球数最小的核。以四个像素块高斯球数 50、100、10、20分配到 2 个核为例均匀切分时 core-0 承担 150 个、core-1 承担 30 个严重倾斜负载均衡后 core-0 承担 100 个、core-1 承担 80 个分配过程tile-1→core-0tile-0→core-1tile-3→core-1tile-2→core-1。两种策略的对比见 切分策略对比图贪心分配过程见 贪心分配示意图。在 rasterizer.py 中负载调度由get_render_schedule在 CPU 侧根据tile_sums每 tile 高斯计数与 vector 核数acl.get_device_capability(0, 1)[0]获取生成lb_sched_tensor再传给排序与渲染算子。对比实验中13 万高斯球、过滤后 67917 个正向与反向耗时收益均约 23%优化方法前向 device 耗时(ms)反向 device 耗时(ms)按块切分5.97315.908负载均衡4.59912.223优化四Precise Intersection 融合算子FlashGaussianBuildMask算子基于 FlashGS 算法。计算流程对每个像素块计算全部高斯球相切矩形边界框顶点坐标判断边界框是否与像素块相交得到掩码mask判断中心点是否在像素块范围内得centerFlag判断像素块边界是否与高斯球投影相交得intersectFlag三掩码做与或操作得到最终映射关系。Tiling 设计将像素块数M作为总任务数分核到各 AICore仅用 Vector 核采用大小核分核策略前 n 个核计算 K 个任务、后blockDim-n个核计算 K-1 个任务见 分核示意。切块方面将 UB 的 85% 用作高斯球搬入/搬出/分块计算区15% 用于临时 Buffer一次切块计算 2.62K 个高斯球见 切块示意。Vector-MTE 流水并行常规核内计算为 MTE2 搬入 → Vector 计算 → MTE3 搬出的线性流程需两次流水同步。优化后当前分块计算完成、MTE3 搬出结果时由于搬出数据与下一分块计算相互独立MTE3 搬出与下一分块计算并行若当前分块是最后一个像素块的映射计算则在计算过程中同步搬入下一分块所需数据使 Vector 核不必等待 MTE2 同步。在 mipnerf360_v2 的 garden 场景上搬出并行收益 7.9%叠加搬入并行再收益 6.8%总体 15.2%。计算优化高斯球属性2D 投影坐标、协方差逆、不透明度在全部像素块计算中不变故将计算拆解为高斯球相关特征计算边界框、不透明度对数与高斯球-像素块映射关系计算两步避免重复搬入。同时基于 AscendC 编译器自动插入依赖同步的特性尽量减少上下文依赖例如边界框宽高公式w ⌊√(2·cov₀₀·power)1⌋、h ⌊√(2·cov₁₁·power)1⌋常规依次计算优化后同步计算2·cov·power、再同步开方加一、最后同步取整减少同步次数。流水并行基础上叠加计算拆解收益 21.1%再叠加计算过程调整收益 4.6%总体 26.7%。高斯球坐标初筛与掩码过滤由于按像素块分核各核上像素块纵坐标范围不完全重叠可在相交计算前先按核的 Y 坐标范围初筛高斯球过滤约 90% 的高斯球使后续相交计算仅需 2~3 次循环Bicycle 场景实测。配合掩码过滤每一轮mask、centerFlag、intersectFlag只保留有效高斯球的属性与索引避免冗余开销。在 mipnerf360_v2 场景上的对比基线为已应用流水并行与计算优化的算法高斯球数量从 20w 到 300w收益 4.7%~30.3%总体约 18.8%。与 GS 算法的端到端对比高斯球数量像素块数量映射关系计算方法渲染前向 device 耗时(ms)20w425GS /Precise Intersection8.955 /6.67140w425GS /Precise Intersection17.499 /13.017100w425GS /Precise Intersection35.957 /22.8807w1107GS /Precise Intersection7.783 /4.86314w1107GS /Precise Intersection15.384 /9.48928w1107GS /Precise Intersection25.311 /15.493可以看到Precise Intersection 在各高斯球规模与像素块规模下均明显优于传统 GS 映射剔除更多无效高斯球从而降低渲染耗时。总结与适用前提本项目通过四类 Ascend C 自定义算子将 3DGS 训推链路球谐解码 → 融合投影 → tile 映射 → 深度排序 → Alpha blending完整移植到昇腾 Atlas A2 SIMD 架构并逐环节优化流水与双缓冲掩盖同步与访存开销、两级剪枝提前终止无效计算、贪心负载均衡消除核间拖尾、精确相交映射与逐级过滤减少冗余计算。上述性能数据均在仓库指定数据集mipnerf360_v2 360_v2 子集、约 13 万高斯球场景与软硬件组合CANN 8.2.RC1、torch 2.1.0 / torch_npu 2.1.0.post12、gcc 10.2、Python 3.8 单卡下测得更换数据集、卡型或软件版本时结果可能不同建议以 ops/ascendc/tests 下的逐算子精度测试为基准进行复现与调参。Citation本样例与优化实践参考以下公开工作article{kerbl3Dgaussians, author{Kerbl, Bernhard and Kopanas, Georgios and Leimk{\u}hler, Thomas and Drettakis, George}, title{3D Gaussian Splatting for Real-Time Radiance Field Rendering}, journal{ACM Transactions on Graphics}, number{4}, volume{42}, month{July}, year{2023} } article{ye2025gsplat, title{gsplat: An open-source library for Gaussian splatting}, author{Ye, Vickie and Li, Ruilong and Kerr, Justin and Turkulainen, Matias and Yi, Brent and Pan, Zhuoyang and Seiskari, Otto and Ye, Jianbo and Hu, Jeffrey and Tancik, Matthew and Angjoo Kanazawa}, journal{Journal of Machine Learning Research}, volume{26}, number{34}, pages{1--17}, year{2025} }此外Precise Intersection 算子基于 FlashGS 算法Feng 等arXiv:2408.07967实现Alpha blending 优化细节可进一步查阅 NPU 3DGS 训推优化实践文档 中的四篇专题。【免费下载链接】cann-recipes-embodied-ai本项目针对具身智能业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-ai创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考