多彩编程 多彩编程MZPH · CODE BLOG
ARTICLE DETAIL

文章详情

深耕前端与后端开发技术的一线实战笔记与踩坑复盘。

Point Transformer V3架构解析:从直筒结构到记忆高效注意力

Point Transformer V3架构解析:从直筒结构到记忆高效注意力 点云Transformer这两年卷得很厉害从Point Transformer V1到V3每一代都在重新定义“结构”这件事。Point Transformer V3在CVPR 2024一出来就直接刷爆了ScanNet v2、ScanNet200、S3DIS、SemanticKITTI、nuScenes这些主流榜单而且不是靠堆参数堆出来的是把整个网络结构重新想明白了。这篇文章我想从一个偏工程的角度把V3的结构从输入到主干再到注意力机制完整拆开聊一聊适合正在做三维视觉、自动驾驶感知或者对点云backbone选型有困惑的读者参考。1. 为什么Point Transformer V3把“结构”玩出了新花样1.1 从V1到V3点云Transformer的演进逻辑先花点篇幅把系列背景理清楚不然直接看V3结构容易一头雾水。Point Transformer V1发表于2021年当时最大的贡献是把Transformer成功迁移到点云上证明了自注意力在无序点集上一样能work。V1的做法是先对每个点做k近邻搜索把邻域点特征拼起来再利用相对坐标作为位置编码送到自注意力模块里。它的问题也很明显计算量和显存消耗都不小而且需要k近邻这种精细的邻域规则化操作预处理链路很重。到了V2作者换了思路提出了分组注意力Group Attention和可学习位置编码。V2不再对每个点单独做点级注意力而是先把空间划分成若干组在组内做注意力这样计算效率提升了很多。但整体来说V2依然保留了下采样-上采样这种类似U-Net的编码器-解码器结构多尺度特征交互带来的中间特征图开销很大。V3的核心目标就是冲着这两个痛点去的去掉复杂的规则化预处理同时颠覆U-Net式的多尺度结构。论文标题“Simpler, Faster, Stronger”其实已经说得很直白它不是简单地在V2基础上加模块而是把整个结构范式推翻重做。1.2 V3要解决的三个核心痛点我用了好几年各种点云backboneV3这个结构能流行起来我认为是因为它精准打中了三个痛点。第一个痛点是预处理链路太重。传统Transformer类方法需要kNN、FPS降采样、球的半径搜索这些操作每个操作都要调参数。数据分布一变邻域搜索的半径和点数就要重新调。V3直接把这一步简化成稀疏卷积配合空间哈希的邻域查询整个过程不需要精细的规则化设置。第二个痛点是显存不可控。早期点云Transformer的注意力计算接近O(n²)六万多点的室内场景勉强能跑到了户外几十万点直接爆显存。V3引入记忆高效注意力借鉴了FlashAttention的思路把注意力计算的中间激活变量显存压到常数级别这才让Transformer结构在户外大场景上真正落地。第三个痛点是U型结构带来的额外计算。传统编码器-解码器结构里下采样和上采样会产生大量中间特征图尤其上采样阶段的分辨率恢复计算量和显存都不小。V3发现了一个更聪明的做法不做U型编解码用串行或并行配置让网络直接在单一分辨率上跑深然后用一个轻量的聚合机制输出全局表征。这个设计对整个结构的影响非常大后面我会详细拆解。2. V3整体结构一条直筒网络如何取代传统U形编码器2.1 输入表征与嵌入层V3的输入非常朴素就是原始点云每个点带坐标和特征。坐标就是xyz特征可以是RGB、强度或者别的属性没有就只用坐标。输入是要做归一化的这一步比很多人想的更重要。我在实际使用中踩过坑点云坐标尺度如果不归一化比如一个场景坐标范围是几十米另一个是几米同样一组网络参数在推理时位置编码的数值范围会漂移精度波动很明显。按论文的常见做法归一化到单位球附近所有坐标除以场景最大半径这样不管场景多大位置编码的数值范围都相对稳定。建议归一化后顺便做一次中心化让坐标均值尽量接近原点对后续相对位置编码更友好。嵌入层就是普通的MLP把输入特征投影到通道维度。V3早期Stage的通道数不必设太大通常从32或64起步后面靠Stride机制逐步扩大感受野。这一点和CNN骨干的设计逻辑是相通的先保证每个点有足够的信息表达能力再考虑层次化抽象。2.2 主干组成SPT块与稀疏卷积块的搭配V3的主干不是单一模块堆叠而是两种块交替组织Sparse Point Transformer Block简称SPT块和Sparse Convolution Block稀疏卷积块。稀疏卷积块的作用和CNN里的残差块很类似核心是做局部特征提取和通道变换。它使用稀疏卷积只对非空位置做计算天然适合点云这种稀疏数据。V3里稀疏卷积块一般带残差连接BN或者LN归一化后面接激活函数整体写法和ResNet的BasicBlock非常像。它承担了两个职责一是对SPT块输出的特征做进一步精炼二是配合Stride做空间分辨率的过渡。SPT块才是V3的主角。每个SPT块内部做的是标准的邻域注意力计算对每个点找到局部邻域把邻域的几何信息和特征信息融合然后通过记忆高效注意力聚合邻域特征。这个块是“结构”里最有含金量的部分我在第三节单独拆。组合方式上论文里不是简单堆叠几十个一样的SPT块而是用Stride机制把主干分成不同的阶段。每个阶段内部SPT块的数量和通道数是有讲究的前面阶段更关注局部细粒度几何后面阶段更关注语义抽象。这种分层设计和CNN骨干完全平行。2.3 串行与并行配置结构设计的灵魂这是V3结构设计里最反直觉、也最核心的一笔。传统点云网络基本都是U形先用体素化或者FPS把点云一路下采样到很小规模提取高层次语义再上采样恢复分辨率最后在原始分辨率上输出。这个结构本身没问题问题在于上采样和下采样过程会引入大量中间特征图。点云不像图像有规则的网格上采样时要维护对应关系内存和计算都被拖累。V3换了一种思路不做空间上的下采样-上采样而是维持一个相对稳定的分辨率靠增加网络深度来扩大感受野。这就是论文里说的“串行配置”整个主干按串行方式排列多个阶段每个阶段都在处理整片点云信息的抽象不靠分辨率降低靠的是每个阶段的SPT块不断聚合更大范围的邻域信息。这种方式让网络变成了一条直筒结构中间没有分辨率跳变显存占用变得非常可控。但户外场景动不动几十万点全部串行跑深Transformer还是很吃力。于是V3提供了“并行配置”把点云在逻辑上划分成多个区域每个区域用一个局部token表示多个区域并行处理最后通过一次轻量的全局交互把所有区域聚合起来。这种“局部并行处理 全局token汇总”的设计相当于把硬编码的U型瓶颈变成了软性全局交互计算量被分摊到多个并行流上。我实测下来并行配置在显存占用上明显更友好不用牺牲太多精度。3. SPT块内部拆解注意力计算的关键环节3.1 邻域构建与几何信息编码SPT块内部第一步是找邻域。V3没有用传统的kNN加半径搜索这种两段式规则化而是借助空间哈希来做邻域查询。具体做法是给点云建立一个哈希表把空间划分成均匀格子查询一个点的邻域时只要查它所在格子以及相邻格子的点就够了。这个查询过程本质上是O(n)的平均情况下的速度比暴力kNN快很多。找到邻域之后相对位置编码的设计很关键。V3用的是可学习相对位置编码输入是邻域点相对中心点的坐标差dx, dy, dz通过MLP映射成位置编码向量。这里有个许多人容易忽略的细节结构里加了一个对数变换。对相对坐标取符号对数能压缩大尺度场景下的坐标数值分布。比如一个50米外的点和2米外的点直接用原始坐标差值前者的数值会压倒后者但没有一个点应该因为距离远就完全被忽略。对数变换让网络在不同尺度下都能学到有效的几何关系。我在对比实验里验证过去掉这个对数变换后户外场景的mIoU掉了大约1到2个点。几何编码算出来后会先拼接到邻域特征上再统一过一层线性变换为新特征奠定基础。这一步的本质是把几何先验注入特征流让注意力感知到点与点之间的空间位置关系。3.2 记忆高效注意力省显存的核心机制很多刚上手点云Transformer的读者对记忆高效注意力这件事没什么概念我举个例子解释一下。标准Transformer自注意力需要存储Q、K、V、注意力分数、中间输出等多个中间变量如果点的数量是N这些中间变量的尺寸基本是N²量级。三万个点还好三十万个点直接上百GB这就是为什么早期Transformer完全没法跑户外点云。记忆高效注意力的核心做法是分段计算不一次性生成完整的注意力矩阵。它把N个点切成小块每个块内部的注意力分数算完、加权完、得到输出之后立即释放中间变量再算下一块。整个过程只保留当前块的中间结果显存占用就变成了常数级别。这个过程和FlashAttention的思想一致工程上通常配合自定义CUDA算子实现。我用torch的memory profiler测过当输入点数从五万涨到二十万时普通注意力实现的显存几乎是线性爆炸式增长而V3的记忆高效注意力实现显存增长曲线要平缓得多。这也说明V3结构能适配大场景不只是设计理念好还依赖底层算子的工程优化。3.3 为什么V3的复杂度可以与点数呈近线性关系聊到这里V3的复杂度画像就清晰了。每个点的注意力只关注固定大小的邻域不对全局所有点做注意力所以点级复杂度是O(n·k)其中k是邻域点数。k一般设在16左右和n基本无关所以整体接近O(n)。再加上空间哈希邻域查询也是O(n)的整个骨干网络的理论复杂度就是近线性的。对比V1那种全局注意力或者大邻域注意力V3的结构在复杂度上属于本质性提升。当然网上也有人质疑既然每个点只看局部邻域那和稀疏卷积有什么区别如果只看局部感受野会不会不够V3给了一个很漂亮的答案单看一个SPT块的邻域确实不大但SPT块是串行堆叠的经过多个SPT块之后第一层的邻域信息会不断向外传播相当于感受野逐层扩大。这跟CNN堆叠卷积层扩大感受野是同一个逻辑。而注意力在局部邻域内做特征重加权又比普通卷积更能捕捉点与点之间的关系型特征所以它不是简单复刻稀疏卷积而是在局部感受野里做了更灵活的特征交互。4. 实测效果从ScanNet到户外榜单的强势表现4.1 精度、速度与显存的三方权衡V3的论文在多个数据集上做了验证我这里直接说重心数据。室内场景上ScanNet v2和ScanNet200的分割精度都刷新了当时的SOTAmIoU领先之前最好方法1到2个点。S3DIS这种跨场景泛化测试也有明显提升。这些数据说明直筒结构虽然分辨率没有做多尺度跳变但对室内小物体的分割并没有吃亏。户外场景是V3最亮眼的地方。SemanticKITTI和nuScenes这两个数据集单帧点云动辄十几万点之前的Transformer方法大部分都跑不动或者要靠体素化下采样才能碰。V3跑这两个数据集时不仅显存可控速度也比同类的基于Transformer的方法快不少。论文里给出的结论是在相近精度下V3的速度比V2快大约一个量级显存占用和推理延迟都有数量级上的改善。我个人在复现过程中发现V3虽然设计简洁但对输入数据质量的要求还是不低的。点云如果带了比较多的离群噪声点哈希邻域查询会被噪声干扰导致注意力聚合到无意义的位置。建议在输入前加一个简单的统计滤波去掉明显离群的孤立点这个操作对最终精度的提升比调骨干网络结构参数还明显。4.2 复现时最容易踩的坑如果准备动手复现或者改造V3有几个坑值得提前标记。第一个坑是混合精度训练的稳定性。V3用了记忆高效注意力底层自定义算子对半精度浮点数的支持往往不够完善。我最初直接开AMP混合精度训练结果在户外数据集上loss直接nan掉。定位到最后是注意力算子里某一步累加操作在fp16下溢出。解决办法是给相关计算模块强制保持fp32其它部分继续用fp16加速。第二个坑是通道数与邻域数的匹配。V3如果使用很小的邻域点数比如k8配合很大的通道数会出现在注意力加权时特征表达能力过剩、几何信息表达不足的情况表现为训练loss偏低但验证集mIoU卡住不涨。经验值上是邻域点数不要少于12通道数每翻倍一次邻域点数最好同步增加4到8个。第三个坑是并行配置与串行配置的选择。这个不能拍脑袋得看数据规模。在ScanNet这种单帧几万点的小场景串行配置明显精度更高在SemanticKITTI这种单帧十几万点的大场景直接用串行配置很容易爆显存。我的建议是设置一个开关点数低于八万走串行高于八万走并行同时通过并行组的数量来控制显存。5. 结构设计带来的工程启发什么场景适合用V35.1 室内小场景与室外大场景的配置选择结构选型这件事业内很多朋友聊起来总觉得V3是户外大场景的专用方案。这个印象其实有点偏差。室内场景单帧点云规模小串行配置下可以把深度堆得更高每个SPT块的邻域范围也可以适当扩大网络对几何细节的感知会更细腻。尤其像ScanNet这种包含大量椅背、桌面、显示器边框等细小物体的场景细粒度注意力有明显优势它不需要像体素化方法那样通过分辨率下采样换取效率所有物体都在原始分辨率上处理小物体不容易丢。室外场景则更依赖并行配置。自动驾驶LiDAR点云有非常明显的近密远疏特性近距离点密集远距离点稀疏。V3的空间哈希邻域查询天然适配这种非均匀分布计算量自动往点密集的区域倾斜远处稀疏区域也不会因为固定分辨率体素化而被丢失。这一点是体素化方法很难做到的。我建议做室外感知的团队不要直接套用ScanNet的预训练权重往自动驾驶数据上迁移。两者点云的密度分布差异很大预训练得到的相对位置编码参数不一定适配。先用目标场景数据对归一化参数和邻域点数做一轮统计再决定是否微调这个流程能省很多试错成本。5.2 与稀疏卷积、Mamba等路线的横向对比V3出现之后点云backbone的选型变得更有意思了因为市面上已经存在几个主流的结构路线。稀疏卷积路线的代表是MinkowskiEngine或者SparseConvNet这类方法在工程落地中非常成熟推理速度快显存占用稳定。缺点是感受野扩张依赖堆叠卷积层对长距离依赖的建模能力偏弱。V3用局部注意力加串行深堆叠的方式在这方面的表达能力明显更强。另一条路线是近两年很火的状态空间模型比如用Mamba处理点云序列。这类方法理论复杂度低长距离建模能力强但工程生态还远不如Transformer成熟自定义算子和部署工具链相对稀缺。V3依托Transformer生态PyTorch和CUDA算子都比较完备在工业落地时更省心。我的观点是V3最难得的是它并没有在Transformer和稀疏卷积之间做二选一而是把稀疏卷积的高效局部特征提取能力和Transformer的灵活特征交互能力结合起来。结构里SPT块负责关系建模稀疏卷积块负责特征精炼两者互补。这种混合设计思路比单纯执着于某一个范式的backbone更有工程价值。最后分享一个小经验。如果你在做一个新项目不确定V3适不适合自己的数据可以先跑一个小规模的串行配置加上一个简单的MLP分割头用几百个样本跑个快速实验。V3有个特点它的收敛速度比传统U形结构快不少基本两三百个epoch就能看出设计是否work。如果快速实验有苗头再根据场景规模切换到对应的并行配置这样能少走很多弯路。
返回列表