折腾了两周,终于在 PyTorch 里把 RandLA-Net 在 Semantic KITTI 上跑通了。说“跑通”不是指 loss 能往下走,而是 val mIoU 稳定出现在预期区间,预测结果可视化不再是一团乱麻。复现这个网络的难点不在网络本身——它的核心结构并不复杂,难的是数据管线、环境版本和一堆只能在实跑中暴露的细节。这篇内容写给正在复现或准备复现 RandLA-Net 的读者,尤其是想在 Semantic KITTI 上做语义分割实验、又不想在踩坑上花两周的人。
我默认你已经知道 RandLA-Net 是干什么的:一个面向大场景点云的语义分割网络,亮点是随机采样替代昂贵的 FPS,再用局部特征聚合把丢掉的几何信息补回来。但“知道原理”和“能复现出结果”之间隔着一整条数据流水线。下面这些内容全部来自我这段时间的实际操作记录,包括环境怎么搭、数据怎么读、网络怎么对、参数怎么调,以及我在每一个环节踩过的坑。
1. 复现之前:先搞懂 RandLA-Net 的思路和 Semantic KITTI 的脾气
1.1 RandLA-Net 到底在解决什么问题
很多人一上来就 clone 仓库开始跑,结果模型跑起来但不收敛,回头才发现根本没理解网络的设计动机。RandLA-Net 的核心矛盾是:点云场景越来越大,传统采样策略成了瓶颈。早期方法常用最远点采样(FPS)来降采样,它能保持点云分布均匀,但复杂度接近 O(N²),在百万点级别的 LiDAR 扫描上完全吃不消。RandLA-Net 的方案很直接——用随机采样,复杂度 O(N),快是快了,但随机采样会丢失局部几何信息,于是论文配套设计了局部特征聚合模块(Local Feature Aggregation)来补偿。
这个模块包含三块:局部空间编码(把相对坐标和点特征拼起来过 MLP)、注意力池化(用 softmax 权重把邻域特征加权求和)、以及空洞残差块(用不同尺度的邻域分组扩大感受野)。理解这几块的作用,你才能判断复现代码里哪些细节可以精简,哪些动不得。比如注意力池化里的 softmax 权重如果直接去掉改成平均池化,模型在小物体上的表现会明显变差,这不是超参问题,是结构问题。
1.2 Semantic KITTI 的数据组织和标签编码,很多人第一步就搞错了
Semantic KITTI 基于 KITTI 里程计数据集,一共 22 个序列,其中只有 00~10 提供了语义标签。常规做法是拿 00~07、09~10 当训练集,08 当验证集,这也是大多数论文和复现仓库的标准划分方式。数据本身不复杂,但格式细节很致命:点云存在 velodyne 目录下,.bin 文件每行四个 float32,分别是 x、y、z、intensity;标签存在 .label 文件里,每个点一个 uint32,低 16 位是语义标签,高 16 位是实例标签。
这里必须强调一下读取顺序,我第一次就栽在这上面:
points = np.fromfile(bin_path, dtype=np.float32).reshape(-1, 4) raw_label = np.fromfile(label_path, dtype=np.uint32).reshape(-1) semantic_raw = (raw_label & np.uint32(0xFFFF)).astype(np.uint16) # 语义标签在低 16 位 instance_raw = (raw_label >> np.uint32(16)).astype(np.uint16) # 实例标签在高 16 位拿到 semantic_raw 之后还不能直接用,原始标签里包含移动的车辆、行人等动态类(比如 moving car、moving person 这类 id),这些类别在训练时应该映射到对应的静态类别上。同时类别 0 是 unlabeled,计算 loss 和 mIoU 时都要忽略。很多复现仓库会提供一个 learning_map 字典,把原始 id 映射到 0~19 的训练 id,这个映射必须保证训练和验证完全一致,否则会出现“训练时用映射后的标签,验证时用原始标签”这种低级错误,最后 mIoU 怎么算都对不上。
1.3 环境选型:PyTorch、CUDA、Anaconda 怎么配才靠谱
复现类项目的第一原则是:先固定环境,再跑代码。RandLA-Net 原版是 TensorFlow 写的,PyTorch 复现的生态比较散,不同仓库依赖的 torch 版本、CUDA 版本差异很大。我建议用 Anaconda 建一个独立环境,不要直接在系统 Python 里装,否则某个依赖升级就会连带搞挂另一个项目。
具体版本上,Python 3.8 或 3.9 基本通吃,PyTorch 我用的是 1.13 + CUDA 11.7,稳定且能覆盖绝大多数复现仓库的要求。如果你的显卡驱动较新、CUDA 主版本是 12.x,装对应版本的 torch 也行,但注意 torch 的 CUDA runtime 和驱动的次版本匹配并不是越新越好,装完后第一件事是验证:
import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))三行输出都正常再继续。另外提一句,在 WSL 里跑同样可行,但 WSL 的 GPU 透传对驱动版本更敏感,如果torch.cuda.is_available()返回 False,先别急着重装 torch,检查驱动和 WSL 内核版本更有效。环境一旦跑通,立刻用conda env export把环境导出保存,这个习惯能救你第二次复现时的命。
2. 数据管线:预处理才是复现 RandLA-Net 最大的坑
2.1 坐标归一化:小决定,大影响
Semantic KITTI 的 LiDAR 点云坐标是实际物理尺度,范围可能是几十米甚至上百米,而网络里 MLP 的初始权重一般在零点几的量级。如果直接把原始坐标喂进去,前几层的输入数值范围巨大,训练必然不稳定,loss 曲线像心电图一样乱跳。所以主流复现都会做归一化,常见两种:一是减均值除以标准差,二是按场景范围缩放到 [-1, 1]。
关键是统计量的来源。正确做法是从训练集算出全局的均值或范围,然后固定下来,训练和验证都用同一套参数。我见过有人对每一帧都独立做归一化,也有验证时重新统计的,前者是数据增强之外的额外随机扰动,后者相当于让模型看到了验证集的分布信息,都会让结果不可复现。intensity 通道的处理也要想清楚:用就用,不用就整体去掉,不要这帧用那帧不用。我自己最后选了“xyz 按全局范围和均值归一化 + intensity 单独做一次 min-max 到 [0,1]”的方案,实测收敛稳定。
2.2 切块、固定点数采样与 KNN 计算
LiDAR 单帧点云数量不固定,而网络输入需要一个固定点数。RandLA-Net 的常见做法是每帧随机采样固定数量的点,比如 40960 个,然后再为每个点算 K 近邻(常见 K=16)。这里有两个容易忽略的细节。
第一,随机采样会让样本之间点数完全一致,但采样结果每 epoch 都不同,相当于一种天然的数据增强。如果想要结果可复现,就必须固定随机种子或者把采样索引缓存下来;反之如果想要更好的泛化,就不要固定。这两种取向没有对错,但你要知道自己用的是哪种。
第二,KNN 索引的计算开销很大。40960 个点、16 个邻居,用 scipy 的 cKDTree 算一帧也要零点几秒,训练 100 个 epoch 就是几百分钟纯浪费。我强烈建议把索引缓存到磁盘,首次计算后保存成 .npy 文件,后续直接加载:
import os import numpy as np from scipy.spatial import cKDTree def get_knn(points, k, tag): cache_path = f"cache/knn_{tag}.npy" if os.path.exists(cache_path): return np.load(cache_path) tree = cKDTree(points) _, idx = tree.query(points, k=k) np.save(cache_path, idx) return idx一个值得利用的性质是:刚体变换(旋转、平移)不改变点与点之间的距离关系,所以 KNN 索引在旋转增强后依然有效,完全不需要重算。但缩放和随机丢弃点会改变邻域关系,这时候索引必须重新生成,切记。
2.3 标签读取与映射的正确姿势
标签这块我在前面提过读取方式,这里再展开说映射。原始标签里有很多动态类别,比如 moving car 是一个独立 id,静态 car 是另一个 id,如果不去映射,模型会被“同一个物体两种标签”搞糊涂。正确的做法是维护一个字典,把原始 id 映射到训练 id,且把 unlabeled 映射成一个专门的忽略类。
实现映射时别用循环,直接向量化:
mapped = np.zeros_like(semantic_raw) for raw_id, train_id in learning_map.items(): mapped[semantic_raw == raw_id] = train_id然后 loss 里设置ignore_index=0(或者你定义的那个忽略类对应的 id)。写完映射后,我强烈建议写一个单元测试:随机抽 5 帧,统计映射前后的类别分布,确认没有出现未处理的原始 id,再确认移动类别正确合并到了静态类别。这一步只要花半小时,能省掉后面排查 mIoU 的十个小时。每次动手调模型前,先用这个小测试确认数据是对的——这也是我这次复现学到的最深刻的一条教训。
3. 网络实现细节:encoder-decoder 与局部特征聚合
3.1 Local Feature Aggregation 的三个子模块怎么实现才不容易错
PyTorch 复现的核心工作就是把论文里的局部特征聚合模块翻译成张量运算。这个模块的输入是当前层的点坐标和特征,输出是聚合后的特征,包含三个子步骤。
首先是局部空间编码。对每个点取 K 个邻居,计算邻居点与中心点的相对坐标,把相对坐标和中心点特征(以及邻居特征)拼起来,过一个共享 MLP。这里最容易出 shape 错误:特征张量的维度一般是 [B, N, C],邻居特征 gather 之后是 [B, N, K, C],相对坐标是 [B, N, K, 3],拼接后的维度是 [B, N, K, C+3] 或 [B, N, K, 2C+3](取决于你拼不拼中心特征),MLP 作用在最后一维上。很多仓库在这里用unsqueeze广播时把维度搞混,前向传播一跑就报错。
其次是注意力池化。对 [B, N, K, C'] 的特征过一层共享 MLP 得到分数,在 K 维度上做 softmax,然后把邻居特征按权重加权求和,得到 [B, N, C']。这一步的 softmax 维度千万别写错——是在邻域 K 那一维,不是批次 B 或通道 C。我调试时习惯在中间打印shape和一行softmax_weights.sum(-1)的统计,权重和应该全为 1,不为 1 说明 softmax 维度写错了。
最后是空洞残差块。用不同膨胀率(dilation rate)的邻居分组并行提取特征再融合,避免单尺度邻域丢细节。残差连接要注意输入输出通道是否一致,如果不一致要用一个 1×1 卷积做投影,否则张量加不起来。这里的坑集中在邻居索引的构造成——膨胀率变了,但索引依然来自同一个 KDTree,只是取邻居的间隔不同,写代码时别重新算一遍 KNN。
3.2 Encoder-Decoder 结构与上采样方式
RandLA-Net 的主体是 5 层编码器,每层先用局部特征聚合提特征,再随机采样把点数减半,通道数逐层翻倍(常见配置从 8 或 16 起步),最深层点数约为输入的 1/32。解码器再用最近邻插值或按距离加权的方式把特征逐步上采样回原始点数,并与对应编码层的输出做 skip connection 拼接。
PyTorch 复现时要注意两个问题。第一,网络对输入点数必须是弹性的。虽然训练时固定采样 40960 点,但验证或推理时你可能想用完整帧,中间层的点数会随之变化,任何把 N 写死的代码(比如x = x[:40960]或硬编码 reshape)都要清理干净。第二,上采样和 skip connection 拼接时,要保证来自解码器的特征和来自编码器的特征点数一致、通道拼接顺序一致(先拼哪个都要统一)。最方便的做法是把采样比例定义成常量,比如每次除以 2,任何一层出了问题都能通过打印每层点数快速定位。
3.3 损失函数与类别不均衡处理
Semantic KITTI 的类别分布极度不均衡,car 和 unlabeled 占了大部分样本,行人和自行车这种小类别占比可能不到 1%。如果直接用普通交叉熵,模型会倾向于把所有点都预测成大类别,val mIoU 看起来还行,实际上细看每一类的 IoU,小类别全是 0。
解决方法是类别加权交叉熵。权重可以从训练集统计类别频率,常见的做法是把每个类别的频率取中位数后,权重大致等于 median_frequency / class_frequency,这样既不会把大类别压得太死,也能把稀有类别抬起来。计算频率时注意只用有效类别(排除 ignore),并且用与训练完全相同的映射后的标签来统计。还有一点:PyTorch 的CrossEntropyLoss自带ignore_index参数,但权重向量必须对应训练 id 的顺序,很多人把权重和原始 id 对应错了,损失函数整体错位却浑然不觉。
4. 训练实测:显存、超参数与收敛问题排查
4.1 显存爆炸:先算账,再调参
RandLA-Net 在训练时最吃显存的地方就是各层的局部特征聚合。一个 [B, N, K, C] 的中间张量,假设 B=2、N=40960、K=16、C=64,单是这一个张量就有 2×40960×16×64×4 字节,约 64 MB,而网络里这样的张量不止一个,再算上反向传播保留的中间梯度,几个 GB 瞬间就没了。这就是为什么很多人一跑训练就 OOM。
遇到 OOM,我建议按这个顺序处理:先把 batch size 降到 1,再降 num_points(40960 降到 24576 或 16384),这两个是主要变量。如果还不够,考虑梯度累积来模拟更大的 batch,或者用 AMP 混合精度。但混合精度要小心 BatchNorm——BN 的均值和方差建议保持 fp32 计算,否则训练过程容易出现诡异的波动。另外,KNN 索引如果每次都在 GPU 上现算,非常耗显存和算力,我直接把索引缓存到 CPU,用torch.from_numpy转成 tensor 后gather特征,实测显存和速度都有改善。
4.2 Loss 不下降、震荡、NaN 的排查顺序
训练中最常见的是三种症状:loss 完全不动、loss 剧烈震荡、loss 直接变成 NaN。按我的经验,排查顺序应该是:先看数据再调参数。第一步确认标签没有 -1 或 NaN——标签里出现 -1 会流入损失函数,直接导致反向传播异常。第二步看学习率,PyTorch 复现里比较常见的初始学习率在 1e-2 到 1e-3 之间,如果你把原版 TF 代码里的 1e-2 直接搬过来,而你的归一化方式和原版不完全一样,很可能剧烈震荡;这时候降到 1e-3 甚至 5e-4,通常立刻稳住。
第三步看梯度,用torch.nn.utils.clip_grad_norm_(model.parameters(), 10.0)限制一下梯度范数,防止个别 batch 产生大梯度把权重炸飞。第四步看是否有个别层输出分布异常,我曾在 debug 时发现归一化把坐标缩到 1e-6 量级,所有 MLP 的输出几乎都是常数,loss 自然降不下去。还有一种隐蔽情况:DataLoader 多进程下 collate 出错导致某些 batch 数据错位,特征是 loss 偶尔跳一个很大的值,这种要靠固定 random seed + 单进程调试才能定位。
4.3 验证 mIoU 的规范操作:别让评估过程骗了你
训练完成后,验证评估要遵循一套固定流程。逐帧读取验证序列(比如 08 序列)的点云,走一遍和训练完全相同的归一化、近邻计算,用训练好的模型输出每个点的 logits,取 argmax 得到预测的训练 id,再通过逆映射还原回原始语义标签,与真实标签逐点对比计算 mIoU。计算时需要忽略 unlabeled 类,否则这个占比较大的类会拉高整体指标,制造一种“表现良好”的假象。
我见过有人用验证 loss 来选最优 epoch,这是不对的。验证 loss 和 mIoU 的走势并不完全一致,尤其类别权重较大时,loss 降低不代表小类别的 IoU 在提高。正确做法是每几个 epoch 做一次完整验证,记录每个类别的 IoU,以验证 mIoU 为标准选模型。如果你发现某个特定类别一直很低(比如 bicycle、person),优先怀疑类别权重算错了,或者该类别在训练集中数量实在太少——前者是 bug,后者才是需要数据增强或重采样解决的问题。
5. 高频问题排查与调试套路
5.1 常见问题速查表
我把这次复现和一些读者反馈中最高频的问题整理成了一张速查表,遇到问题先对着表查一遍,能省不少时间。
| 现象 | 可能原因 | 排查手段 | 解决办法 |
|---|---|---|---|
| 前向传播 shape 报错 | gather/unsqueeze 维度写错 | 打印每层 tensor shape | 统一 [B,N,K,C] 约定,加 shape assert |
| 训练 OOM | num_points 或 batch 太大 | 用 max_memory_allocated 定位 | 降到 batch=1、点数 24576 以下 |
| loss 为 NaN | 学习率过高 / 标签含 -1 | 检查标签分布与梯度范数 | 降 lr、标签映射加校验 |
| loss 不降 | 坐标未归一化 / 权重错 | 检查输入分布与 loss 曲线 | 统一归一化,修复类别权重 |
| 训练很快但 mIoU 上不去 | 小类别被忽略 | 看 per-class IoU | 修权重、开数据增强 |
| 验证 mIoU 与预期差很多 | 训练/验证标签映射不一致 | 打印两边标签唯一值集合 | 共用同一 learning_map |
| DataLoader 极慢 | 每 epoch 重算 KNN | 观察 CPU 占用 | 缓存索引到磁盘 |
| 结果不可复现 | 随机种子未固定 / 增强随机 | 固定 seed 并记录 | 设 torch.manual_seed + numpy seed |
5.2 最小复现与调试三板斧
碰到疑难 bug,我的办法是先搭一个最小可复现的 debug 配置:只用 1~2 个训练序列,一个验证序列,num_points 降到 4096,batch size 为 1,epochs 设 2。用这个配置先把 forward + backward 完整跑通,再逐步放大到完整配置。调试过程中有三板斧非常管用。
第一斧是打印 shape。每一层模块都打印输出 shape,和设计值对比,大部分维度错误一眼就能看出。第二斧是固定随机种子,包括torch.manual_seed、np.random.seed、DataLoader 的generator,这样可以确保同样的代码跑两次结果一致,否则你永远分不清是代码 bug 还是随机性。第三斧是可视化。把验证帧的预测结果按类别上色存成点云文件,用可视化工具打开看,如果预测结果在类别边界出现大面积混叠,说明特征聚合或上采样有问题;如果只是局部噪点,多半是标签映射或后处理的问题。可视化是最直观的调试信号,比盯着一堆指标猜高效得多。
5.3 复现项目的版本管理与实验记录
复现类项目最怕的是“记不清改了什么”。我强烈建议从一开始就建立简单的实验纪律:代码用 git 管理,每改一个关键点提交一次;环境用conda env export > environment.yml固定;每次实验把配置文件(num_points、K、lr、epochs、种子)写进实验日志目录;指标统一记录到一个 CSV 或直接用 TensorBoard。这样哪怕一周后回头查,也能准确知道哪个改动带来了几个点的 mIoU 提升。
还有个实际教训:网上很多 PyTorch 复现仓库质量参差不齐,有些简化了模块实现,有些偷偷改了损失函数,有些数据管线根本跑不通。我的做法是先挑选一个 star 数高、活跃度高的仓库作为基线,跑通后再逐模块对照论文代码审查一遍,确认每个关键组件没有被“朴素化”处理。用两周时间复核一个基线,比用四周时间在一个有隐藏 bug 的仓库上打转要划算得多。
6. 收尾前再分享几个真实体会
这次复现 RandLA-Net 最大的收获不是终于看到 mIoU 涨上去,而是学会了把问题拆成环境、数据、模型、训练四个层面,逐个击破且不互相干扰。我个人实际体验中最值钱的改动有两个:第一是 KNN 索引缓存,训练速度快了三倍以上,这个优化比任何调参都立竿见影;第二是标签映射的单元测试,它让我排查 mIoU 的时间从以天计缩短到以小时计。
另外有个心态上的建议:复现不是终点,只是起点。一旦你把基线跑通并稳定复现出论文量级的结果,接下来才真正进入有价值的工作——替换局部特征聚合模块、实验不同的采样策略、把网络嫁接到自己的数据上。这些后续实验都要靠一套干净、可靠、可复现的代码基础才能支撑起来。
如果你是刚起步复现这个网络,最后再认真提醒一句:先花一天时间把数据和标签彻底捋顺,再动手碰网络。数据管线的坑一旦埋下,会在后面每一次训练和评估中反复干扰你,而它本身修复起来往往只需要半天。先把地基打牢,上面的一切才能稳固。