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

文章详情

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

3D点云分割中注意力机制实战:从模块选型到PointNet++集成,mIoU提升5个点

3D点云分割中注意力机制实战:从模块选型到PointNet++集成,mIoU提升5个点 简介本资源面向三维感知与深度学习方向的研究者与开发者聚焦通过注意力机制提升3D点云语义分割性能这一热点问题适合具备一定深度学习基础、希望深入理解点云特征学习与工程落地的中高级学习者。压缩包共195个文件以89个Python源码与98个编译缓存文件为主辅以说明文档、结果可视化图片与README整体约1.95MB结构紧凑便于快速复现实验。项目围绕空间注意力与通道注意力两条主线展开结合SPVCNN等体素与点云混合网络在SemanticKITTI、Street3D等数据集上给出分割结果与可视化对比源码完整覆盖模型定义、层实现与训练推理流程。已有231人学习下载读者可据此掌握注意力模块的嵌入方式、点云非结构化数据的处理思路以及分割精度提升的验证方法为自动驾驶、机器人导航等场景的算法研究提供可复用的实践参考。1. 3D点云分割遇上注意力机制为什么加对了模块mIoU 能涨 5 个点做 3D 点云语义分割的同行大概都有过这种体验PointNet 跑在 S3DIS 上mIoU 卡在 55% 上下调学习率、换数据增强、加训练轮数折腾一圈涨不到 1 个点。问题往往不在训练策略而在特征聚合方式——点云是无序、稀疏、密度不均的集合卷积核那套「局部窗口加权求和」在 3D 里天然吃亏远处点和近处点对同一个物体的贡献被同等对待边界和细小结构直接糊掉。注意力机制恰好补这个短板它让网络自己学「哪些点该多看两眼」。自注意力建模长程依赖通道注意力重新标定特征响应两者叠在分割头前面S3DIS 上 mIoU 从 55% 推到 60% 以上是常见结果。这篇笔记面向已经跑通过 PointNet 或类似基线、想在不推翻主干的前提下把精度再抬一截的工程师从模块选型讲到代码落地和踩坑源码结构也会一并说清。2. 点云注意力模块怎么选从 Self-Attention 到通道注意力的取舍2.1 点云上做注意力的三个特殊约束在图像里加注意力输入是规整的 H×W 网格位置编码天然存在。点云不一样三个约束直接决定模块能不能用第一置换不变性。点云输入顺序变了输出必须不变。任何依赖绝对位置索引的操作比如按 index 取邻居都会破坏这个性质注意力权重的计算必须只依赖特征本身或相对几何关系。第二密度不均。室内扫描里桌面点密集、墙面点稀疏同一邻域半径内点数可能差一个数量级。如果注意力直接用点对点的相似度密集区域会主导梯度稀疏区域的特征学不动。常见做法是在邻域内先做一次局部聚合比如 max pooling 或加权求和再在聚合特征上算注意力。第三计算量。原始自注意力的复杂度是 O(N²)N 是点数。S3DIS 一个房间几万个点直接全局自注意力显存扛不住。所以实际方案基本都是「局部注意力 全局通道注意力」的组合而不是纯 Transformer。2.2 自注意力、通道注意力、混合注意力的适用场景选模块之前先明确你的瓶颈在哪。如果模型对「同一物体被切成两段」这类长程关联问题处理不好——比如椅子腿和椅背被分成两个类别——那是空间建模不足优先加自注意力。如果模型对「某些通道响应过强、压制了其他特征」敏感——表现为某些类别 mIoU 特别低——那是通道维度的问题加通道注意力更直接。模块类型代表结构计算复杂度适合场景在点云上的典型增益自注意力Transformer block、PCTO(N²) 或 O(N·k)长程依赖、部件关联S3DIS mIoU 2~4通道注意力SE、ECAO(C²) 或 O(C)通道冗余、类别不均衡S3DIS mIoU 1~2混合注意力CBAM 变体、Point TransformerO(N·k C²)精度优先、算力充足S3DIS mIoU 3~5我一般会先加通道注意力因为它改动最小、几乎不增加显存跑一轮看哪些类涨了。如果边界类board、clutter还是不行再叠局部自注意力。一上来就上全局 Transformer训练不稳定不说调参成本也高。2.3 把注意力模块挂到 PointNet 的哪个位置位置比模块本身更影响效果。PointNet 的 SASet Abstraction层做下采样和局部聚合FPFeature Propagation层做上采样和特征传播。注意力加在不同位置作用完全不同加在 SA 层之后对下采样后的稀疏特征做通道重标定影响全局语义。适合通道注意力。加在 FP 层之前对传播前的特征做空间注意力影响边界恢复。适合自注意力。加在分割头之前对所有点特征做最后一次融合影响最终分类。两种都可以但要注意点数已经恢复到原始规模自注意力开销大。我的习惯是在每个 FP 层的 skip connection 汇合之后、上采样之前插一个通道注意力在最后一个 FP 层之后插一个局部自注意力。这样既有全局通道校准又有局部空间细化显存增加控制在 15% 以内。2.4 一个最小可用的通道注意力实现下面这段代码可以直接插到 PointNet 的 FP 模块里。输入是 (B, N, C) 的点特征输出同形状。import torch import torch.nn as nn class PointChannelAttention(nn.Module): 点云通道注意力对每个点的特征通道做重标定 def __init__(self, channels, reduction8): super().__init__() # 全局平均池化到 (B, C)再经过 bottleneck 学习通道权重 self.mlp nn.Sequential( nn.Linear(channels, channels // reduction, biasFalse), nn.ReLU(inplaceTrue), nn.Linear(channels // reduction, channels, biasFalse), nn.Sigmoid() ) def forward(self, x): # x: (B, N, C) # 对点数维度做平均得到每个通道的全局响应 w x.mean(dim1) # (B, C) w self.mlp(w) # (B, C) # 广播回每个点逐通道相乘 return x * w.unsqueeze(1) # (B, N, C)逻辑说明x.mean(dim1)把每个通道在所有点上的响应压成一个标量这一步等价于 SE 的 squeeze。reduction8控制 bottleneck 的压缩比点云特征通道通常 128~512压到 16~64 足够。Sigmoid保证权重在 0~1 之间不会翻转特征符号。参数说明channels必须和输入特征维度一致插在 FP 层时通常是 128 或 256。reduction不建议小于 4否则 bottleneck 太窄通道间关系学不出来也不建议大于 16否则参数量上去但增益不明显。如果显存紧张可以把nn.Linear换成 1D 卷积效果基本一致。2.5 局部自注意力的轻量实现与邻域选择自注意力在点云上必须限制在局部邻域否则 O(N²) 直接爆显存。常见做法是用 KNN 找每个点的 k 个最近邻在邻域内算注意力。class LocalSelfAttention(nn.Module): 局部自注意力在 KNN 邻域内计算注意力权重 def __init__(self, channels, k16): super().__init__() self.k k self.qkv nn.Linear(channels, channels * 3, biasFalse) self.out nn.Linear(channels, channels, biasFalse) self.scale channels ** -0.5 def forward(self, x, knn_idx): # x: (B, N, C), knn_idx: (B, N, k) B, N, C x.shape qkv self.qkv(x).reshape(B, N, 3, C) q, k, v qkv.unbind(dim2) # 各 (B, N, C) # 用 knn_idx 把邻居特征 gather 出来 k_nb torch.gather(k, 1, knn_idx.reshape(B, -1, 1).expand(-1, -1, C)) v_nb torch.gather(v, 1, knn_idx.reshape(B, -1, 1).expand(-1, -1, C)) k_nb k_nb.reshape(B, N, self.k, C) v_nb v_nb.reshape(B, N, self.k, C) # 注意力权重: (B, N, k) attn (q.unsqueeze(2) * k_nb).sum(dim-1) * self.scale attn attn.softmax(dim-1) # 加权聚合邻居特征 out (attn.unsqueeze(-1) * v_nb).sum(dim2) # (B, N, C) return self.out(out)逻辑说明qkv把输入映射成 query、key、value 三组。torch.gather按 KNN 索引取出每个点的邻居 key 和 value。注意力权重是 query 和邻居 key 的点积softmax 归一化后对邻居 value 加权求和。整个过程只在 k 个邻居内计算复杂度 O(N·k·C)。参数说明k16是室内点云的常用值太小k8感受野不够太大k32显存和计算量上升明显。knn_idx需要在数据预处理或 forward 里用torch.topk或scipy.spatial.cKDTree算好注意要排除自身点。如果点云密度变化大可以考虑用球查询替代 KNN保证邻域物理半径一致。3. 从零跑通数据准备、模块插入与训练配置3.1 S3DIS 数据预处理与 KNN 索引构建S3DIS 原始数据是每个房间一个 .txt每行 6 列xyz rgb外加一个 label 列。PointNet 的标准预处理是把房间切成 1m×1m 的 block每个 block 采样 4096 个点。这一步直接影响后续 KNN 的质量。# 目录结构约定 # data/s3dis/Area_1/office_1/office_1.txt # 每行: x y z r g b label python prepare_s3dis.py \ --data_root ./data/s3dis \ --output_root ./data/s3dis_blocks \ --block_size 1.0 \ --num_points 4096 \ --test_area 5block_size1.0是 S3DIS 的惯例对应房间的物理尺度。num_points4096是 PointNet 的默认输入点数如果你的显存够可以提到 8192但 KNN 的 k 要相应调整。test_area5表示用 Area 5 做测试这是标准协议方便和论文对比。预处理完每个 block 存成 .npy包含 points (4096, 6) 和 labels (4096,)。KNN 索引不建议在预处理阶段算因为数据增强旋转、缩放会改变几何关系。我一般在 Dataset 的__getitem__里用sklearn.neighbors.NearestNeighbors现算或者用torch_cluster.knn在 GPU 上算后者快很多。3.2 把注意力模块插进 PointNet 的代码改动假设你用的是 OpenPCDet 或自己写的 PointNetFP 模块的 forward 通常长这样# 原始 FP 模块 forward简化 def forward(self, xyz1, xyz2, feat1, feat2): # xyz1: 上采样目标点坐标, xyz2: 下采样点坐标 # feat1: 上采样目标点特征, feat2: 下采样点特征 dist, idx three_nn(xyz1, xyz2) dist_recip 1.0 / (dist 1e-8) norm torch.sum(dist_recip, dim2, keepdimTrue) weight dist_recip / norm interpolated_feat three_interpolate(feat2, idx, weight) new_feat torch.cat([interpolated_feat, feat1], dim-1) new_feat self.mlp(new_feat) return new_feat插入通道注意力的改动很小# 改动后 def __init__(self, ...): ... self.channel_attn PointChannelAttention(channels128, reduction8) def forward(self, xyz1, xyz2, feat1, feat2): ... new_feat self.mlp(new_feat) new_feat self.channel_attn(new_feat) # 插入位置MLP 之后 return new_feat注意channels要和self.mlp的输出维度一致。如果 FP 层有多级每一级都可以插但建议只在最后两级插前面层特征还不够语义化注意力学不到有意义的东西。3.3 训练超参设置与显存占用对比加注意力之后学习率和 batch size 需要微调。注意力模块的参数量不大但梯度传播路径变长学习率太高容易震荡。配置项基线 PointNet加通道注意力加局部自注意力初始学习率0.0010.0010.0005batch size16168显存占用4096点6.2 GB6.8 GB9.5 GB每 epoch 时间85s92s140sS3DIS mIoU55.357.159.8局部自注意力的显存涨得最明显因为 KNN 的 gather 操作会产生中间张量。如果显存不够两个办法把 k 从 16 降到 8或者只在最后一个 FP 层加自注意力。学习率降到 0.0005 是为了配合更长的梯度路径用 cosine schedule 比 step schedule 更稳。3.4 训练命令与日志观察python train_s3dis.py \ --model pointnet2_attn \ --batch_size 16 \ --lr 0.001 \ --epochs 100 \ --k 16 \ --attn_type channel \ --log_dir ./logs/attn_channel \ --test_area 5跑起来之后重点看两个指标训练 loss 是否比基线下降更平滑验证集 mIoU 在第几个 epoch 开始超过基线。通道注意力通常在第 20~30 epoch 见效自注意力要到 40 epoch 以后。如果 50 epoch 还没动静先检查注意力模块的初始化——Sigmoid输出接近 0.5 时相当于没加如果初始化让权重全接近 0 或 1等于把特征掐死了。可以把最后一层nn.Linear的 bias 初始化为 0weight 用小方差初始化。4. 避坑与排查注意力模块不涨点的五个血泪教训4.1 现象加了注意力mIoU 反而掉了 2 个点原因注意力模块插在了 SA 层之后、FP 层之前但 SA 层的输出特征已经经过 max pooling空间信息损失严重通道注意力在这种特征上学到的权重是噪声。解决把注意力移到 FP 层的 skip connection 汇合之后。FP 层的特征同时包含上采样的语义信息和原始点的几何信息注意力有足够的信息去判断哪些通道重要。4.2 现象训练 loss 震荡验证集 mIoU 忽高忽低原因自注意力的scale因子设错了。scale channels ** -0.5是 Transformer 的标准做法但点云特征经过 KNN gather 之后key 的方差和图像不一样直接用这个 scale 会导致 softmax 输出过于尖锐或过于平滑。解决把scale改成可学习参数初始化为channels ** -0.5让网络自己调。或者用LayerNorm对 q 和 k 做归一化再算点积这样 scale 固定为 1 也不会出问题。4.3 现象KNN 索引算出来全是自身点原因用sklearn.neighbors.NearestNeighbors时忘了排除第一个邻居。KNN 返回的最近点永远是自身距离为 0如果不跳过注意力就变成「自己看自己」等于没加。解决kneighbors返回的 indices 取[:, 1:k1]跳过第一列。或者用torch_cluster.knn时设kk1然后取后 k 个。这个坑很隐蔽因为 loss 不会报错只是不涨点。4.4 现象显存溢出batch size 降到 2 还是 OOM原因局部自注意力在 FP 层做 KNN 时点数已经恢复到原始规模比如 4096gather产生的中间张量是 (B, N, k, C)B8、N4096、k16、C128 时单个张量就是 8×4096×16×128×4 bytes ≈ 268 MB加上反向传播的梯度显存直接翻倍。解决两个方向。一是把自注意力放在 SA 层之后那里点数只有 1024 或 256显存压力小很多。二是用torch.utils.checkpoint对注意力模块做梯度检查点用时间换显存代价是训练速度慢 20% 左右。4.5 现象某些类别 mIoU 涨了另一些掉了原因通道注意力的reduction太小bottleneck 把一些类别的特征通道压没了。比如 S3DIS 里column和wall的特征通道有重叠压缩比太大时两者共享的通道被抑制一个涨一个掉。解决把reduction从 8 调到 4 或 2减少信息损失。或者改用 ECAEfficient Channel Attention它用 1D 卷积代替全连接不做降维参数量更少且不损失通道信息。如果类别不均衡严重可以在 loss 里给稀有类加权配合注意力一起用。5. 进阶技巧用注意力热力图验证模块是否真的学到了东西5.1 可视化注意力权重定位失效模块训练完模型不要只看 mIoU 数字。把通道注意力的权重w和自注意力的attn导出来叠在点云上做可视化能直接看出模块有没有学到有意义的东西。def visualize_channel_attention(model, sample): 导出通道注意力权重并排序 model.eval() with torch.no_grad(): x sample[points].cuda() # (1, N, 6) # 前向到第一个 FP 层取出注意力权重 feat model.sa_layers(x) for fp in model.fp_layers: feat fp(feat) if hasattr(fp, channel_attn): w fp.channel_attn.mlp(feat.mean(dim1)) # (1, C) # 打印权重最高和最低的通道 topk torch.topk(w, 10) print(Top-10 channels:, topk.indices.tolist()) print(Top-10 weights:, topk.values.tolist())如果权重分布接近均匀所有通道都在 0.5 附近说明注意力没学到东西可能是学习率太低或模块位置不对。如果权重两极分化严重少数通道接近 1其余接近 0说明 bottleneck 太窄需要调大reduction。5.2 用消融实验确定模块组合不要一次加多个注意力模块消融实验是必须的。我一般跑四组实验组通道注意力局部自注意力S3DIS mIoU结论baseline无无55.3对照AFP 层后无57.1通道注意力有效B无最后 FP 层58.4自注意力更有效CFP 层后最后 FP 层59.8两者互补如果算力有限优先跑 A 和 B确定哪个方向对你的数据更有效再决定要不要叠 C。注意每组实验除了注意力模块其他超参必须完全一致否则对比没意义。5.3 迁移到其他点云分割任务时的调整这套方案在 S3DIS 上验证过迁移到 SemanticKITTI室外自动驾驶点云或 ShapeNet部件分割时需要调整。室外点云密度更低、距离更远KNN 的 k 要加大到 32 甚至 64否则邻域内点数不够。部件分割的点数少每个形状 2048 点自注意力的全局版本反而可行不需要限制在局部。我自己的习惯是新数据集先跑基线看哪些类 mIoU 最低再决定加哪种注意力。不要因为某个模块在 S3DIS 上涨点就无脑迁移点云分布差异比图像大得多。另外注意力模块的初始化很关键我一般会把通道注意力的最后一层 bias 初始化为 0让初始权重接近 0.5相当于恒等映射训练初期不会破坏预训练特征。希望帮到你。本文还有配套的精品资源点击获取
返回列表