在百舸上使用 SparseBEV++ 加速镜像训练 3D 目标检测模型
概述
SparseBEV 是一种纯视觉（camera-only）3D 目标检测方法：以 ResNet-50 + FPN 提取 6 路相机图像的多尺度特征，通过多尺度多视图可变形采样（MSMV sampling）在 3D 参考点上做稀疏采样，再由稀疏查询解码器（self-attention + temporally-aware attention + adaptive mixing）迭代回归 3D 框。SparseBEV++ 为 SparseBEV 在 nuScenes 数据集上的训练加速优化版本，优化以 RocketX（零侵入训练加速库，以 wheel 形式安装进镜像、通过环境变量开关生效）承载；SparseBEV 源码侧仅保留少量保证源码可跑通的正确性/兼容性修复及多机扩展所需的最小改动。不设任何开关时，训练行为与上游基线一致。
本加速镜像已内置完整训练栈与优化后的训练脚本，适配 RTX Pro 6000（Blackwell，sm_120），单机多卡与多机扩展训练均支持（多机经 RoCE v2 / RDMA 实测拉通至 2 节点 × 8 卡）。本示例以 nuScenes 数据集为例，介绍在百舸平台上完成训练全流程。
适用场景
- 在百舸平台训练 SparseBEV 类纯视觉 3D 目标检测模型，用于自动驾驶感知（nuScenes 等）任务
- 需要单机多卡训练，或多机扩展训练以获得更高吞吐的场景
- 需要断点续训、性能参考与常见问题指引的客户
环境与资源要求
| 项 | 内容 |
|---|---|
| 适配芯片 | 支持规格：aihc.lgn7ib |
| 源码 | 已内嵌于镜像 /root/workspace/SparseBEV（含正确性修复 + 多机改动） |
| 镜像环境 | 自带训练栈（python 3.12 + mmcv 1.7.2 / mmdet 2.28.2 / mmdet3d 1.0.0rc6），开箱即用，无需 venv / conda |
| 工作目录 | /root/workspace/SparseBEV |
| 数据集 | nuScenes（samples/、sweeps/、maps/、sparsebev/ 元信息 pkl），不打入镜像，运行时挂载 |
| 预训练权重 | Cascade Mask R-CNN R50-FPN（nuImages 20e）骨干权重，不打入镜像，运行时挂载 |
在百舸平台准备训练环境
- 登录百舸平台，创建或进入开发机，选择 RTX Pro 6000 资源规格与上述加速镜像
- 打开开发机终端（或 WebIDE 终端），进入工作目录
/root/workspace/SparseBEV - 按本章节顺序准备数据集与权重，然后启动训练
准备数据集与权重
数据集与预训练权重不打入镜像，运行时需挂载到容器内或通过环境变量指定路径。
数据集
nuScenes 数据集需包含以下目录结构（训练/验证元信息 pkl 为含时序的版本）：
1data/nuscenes/
2├── samples/ # 6 路相机图像（CAM_FRONT / CAM_FRONT_LEFT / CAM_FRONT_RIGHT / CAM_BACK / CAM_BACK_LEFT / CAM_BACK_RIGHT）
3├── sweeps/ # 历史帧图像（时序输入）
4├── maps/ # 地图（BEV map 分支用）
5└── sparsebev/
6 ├── nuscenes_infos_train_sweep_temporal.pkl # 训练元信息（含时序）
7 └── nuscenes_infos_val_sweep_temporal.pkl # 验证元信息（含时序）- 可通过 BOS 获取：
bos://ai-model-yq/tianyaolin/sparsebev/datasets/nuscenes/（示例路径，实际以交付信息为准） - 百舸训练任务等不能
-v挂载的场景，设置DATA_PATH=<数据集根目录>即可（见「启动训练」） - 若使用自定义数据，需按同名字段生成对应 pkl
预训练权重
load_from 指向 COCO 预训练的 Cascade Mask R-CNN R50-FPN（nuImages 20e）骨干权重：
1cascade_mask_rcnn_r50_fpn_coco-20e_20e_nuim_20201009_124951-40963960.pth- 可通过 BOS 获取：
bos://ai-model-yq/tianyaolin/sparsebev/pretrained_weights/cascade_mask_rcnn_r50_fpn_coco-20e_20e_nuim_20201009_124951-40963960.pth（示例路径，实际以交付信息为准） - 配置中的
revise_keys = [('backbone', 'img_backbone')]负责把预训练权重里的backbone.*映射到本模型的img_backbone.*
Docker 启动（挂载方式）
镜像内提供 run_docker.sh 运行脚本，自动挂载数据集与预训练权重：
1# 进入交互 shell
2./run_docker.sh
3
4# 直接启动训练（以单机 8 卡、每卡 bs35、60 step 为例）
5./run_docker.sh STEPS=60 GPUS=8 PER_GPU_BATCH=35 bash run_rocketx_verify.sh训练产物默认写在容器内 outputs/（随容器销毁）。需要持久化或做多机共享时，额外挂载卷并设置 OUTPUT_ROOT，例如：
1DOCKER_ARGS="-v /pfs/.../outputs:/verify/outputs" OUTPUT_ROOT=/verify/outputs ./run_docker.sh ...参数说明
训练参数全部通过环境变量传入启动脚本 run_rocketx_verify.sh，无需修改源码或配置文件。
拓扑与启动参数（torchrun）
| 参数（环境变量） | 默认值 | 说明 |
|---|---|---|
NODE_RANK |
0（兼容平台注入的 RANK） |
当前节点编号，从 0 到 NNODES-1，多机每台必须唯一 |
NNODES |
1（兼容平台注入的 WORLD_SIZE） |
参与训练的节点数；所有节点必须使用一致配置，并正确设置 MASTER_ADDR |
GPUS_PER_NODE |
自动探测 nvidia-smi -L 行数 |
每个节点启动的进程数，一进程一卡 |
MASTER_ADDR |
127.0.0.1 |
rank 0 节点地址。NNODES>1 且为 127.0.0.1 时脚本直接报错退出，避免静默死锁 |
MASTER_PORT |
29500 |
torchrun rendezvous 端口，所有节点必须一致 |
TIMEOUT_S |
18000 |
rendezvous 预算（秒），多机冷启动 + 数据加载不易超时 |
训练超参
| 参数（环境变量） | 默认值 | 说明 |
|---|---|---|
PER_GPU_BATCH |
35 |
每张 GPU 的 micro-batch。显存是主要约束：95.6 GiB 卡上为接近上限的配置，调大需关注显存 |
NUM_EPOCHS |
1 |
训练 epoch 数，总步数 = NUM_EPOCHS × STEPS |
STEPS |
28130 // GLOBAL_BATCH |
每个 epoch 的 iteration 数，随全局 batch 自动重算 |
CONFIG |
configs/r50_nuimg_704x256_memory_8x256.py |
训练配置文件（输入 704×256，解码器 8 层、256 维、900 查询、含 DN 分组） |
RUN_NAME |
rocketx_<月日_时分秒> |
本次运行的名称，决定权重与日志落盘目录 |
WORKERS_PER_GPU |
8 |
每个训练进程的 DataLoader worker 数 |
LOG_INTERVAL |
1 |
日志打印间隔，调大可减少日志刷新开销（适合性能测试） |
OUTPUT_ROOT |
outputs |
训练产物输出根目录。多机必须指向所有节点同路径可见的共享 FS |
DATA_PATH |
空 | 数据集根目录。非空时脚本把 data/nuscenes 软链到该目录（用于不能 -v 挂载的场景） |
全局 batch size = NNODES × GPUS_PER_NODE × PER_GPU_BATCH × 1（本配置梯度累积步数为 1）。
权重保存与断点续训参数
| 参数（环境变量） | 默认值 | 说明 |
|---|---|---|
CKPT_SAVE_INTERVAL |
STEPS × 10（按 iteration） |
权重保存间隔，设 100000 等价于关闭中间保存（最后一步仍会保存） |
CKPT_MAX_KEEP |
1 |
保留的权重个数（-1 = 全部保留） |
EVAL_INTERVAL |
STEPS × 10（按 iteration） |
验证间隔，设 0 关闭训练中评测 |
RESUME_FROM |
空 | 断点续训入口，从该权重继续训练，work_dir 取该权重所在目录 |
LOAD_FROM |
空 | 覆盖 config 中的 load_from（只加载模型权重，不恢复 optimizer/iter） |
启动脚本会显式覆盖配置文件中的
checkpoint_config.interval/max_keep_ckpts/eval_config.interval（按 config 内置 batch_size 16 在配置解析期算出，--override batch_size不会重算），避免短跑时「一个中间权重都不存」。
优化配置接口
优化项通过环境变量开关启用（1）/关闭（0）。总开关是 OPT：OPT=1（默认）打开全部性能优化，OPT=0 关闭全部优化（纯基线，用于 A/B 对照）。所有优化项开关默认关，不设开关时训练行为与基线一致。
优化开关一览（run_rocketx_verify.sh 中每项均为 ${VAR:-1}，显式传入的值优先，支持「只关其中一项、其余全开」）：
1export OPT=1 # 1 = 全栈优化（默认）；0 = 纯基线（A/B 对照）
2export ROCKETX_SKIP_VALIDATION=1 # 只跳过云实例 license 校验，不影响任何优化项
3export ROCKETX_VERBOSE=1 # 打印各项命中状态表（首次接入建议打开）
4
5# P0：DN 分桶（去噪查询按 known_num 分桶，削 pad 极值）
6export ENABLE_ROCKETX_SPARSEBEV_DN_BUCKET=1
7export ROCKETX_SPARSEBEV_DN_BUCKET_K=4 # 桶数，默认 4
8
9# P5-a：通用优化旋钮（8 项）
10export ENABLE_ROCKETX_FUSED_OPTIMIZER=1 # AdamW → fused AdamW
11export ENABLE_ROCKETX_GRADIENT_AS_BUCKET_VIEW=1 # DDP 梯度用 bucket view，省一次 grad copy
12export ENABLE_ROCKETX_NUMA_BINDING=1 # 首次 set_device 时绑核，降跨 NUMA 访存
13export ENABLE_ROCKETX_TF32=1 # matmul/cudnn 允许 TF32（引入可接受的数值差异）
14export ENABLE_ROCKETX_TRANSPARENT_HUGEPAGE=1 # 透明大页（需内核 THP 权限）
15export ENABLE_ROCKETX_GC_OPT=1 # gc.disable()，规避周期性 GC 停顿
16export ENABLE_ROCKETX_TENSORBOARD=1 # TensorBoard 记录旁路
17export ENABLE_ROCKETX_PIN_MEMORY=1 # pin_memory：H2D 搬运移出主线程（注意：会改变数据序列，见「常见问题」）
18
19# P5-b：MSMV sampling 反向 CUDA 算子（warp-shuffle 归约）
20export ENABLE_ROCKETX_MSMV_SAMPLING=1
21
22# P6-a / P6-b：按模块关闭激活重计算（显存换反向重算时间）
23export ENABLE_ROCKETX_SPARSEBEV_DISABLE_BACKBONE_CP=1 # 关闭 img_backbone 激活重计算
24export ENABLE_ROCKETX_SPARSEBEV_DECODER_CP=1 # 按模块关闭 decoder 激活重计算
25export SPARSEBEV_CP_OFF=sa,ss # 关闭的 decoder 子模块键 ∈ {sa,ta,ss,am}
26
27# P7：匈牙利匹配向量化 + scipy 并行
28export ENABLE_ROCKETX_SPARSEBEV_HUNGARIAN=1
29export SPARSEBEV_HUNGARIAN_PARALLEL=8 # scipy 并行线程数，设 0 则只向量化不并行
30
31export WANDB_MODE=offline # 默认离线，避免无外网时 wandb 阻塞实现说明：OPT 是启动脚本自身的 shell 分支开关（OPT=1 时逐项 export 各 ENABLE_ROCKETX_*，OPT=0 时全部不设），RocketX 的 is_enabled() 据此判断是否装补丁；变量只向下传给子进程，不影响登录 shell 或其它任务。
启动训练
启动脚本 run_rocketx_verify.sh 已适配单机与多机场景，单机/多机共用同一脚本，仅通过环境变量区分。
单机训练（8 卡）
1# 基线（全关，A/B 对照）
2OPT=0 bash run_rocketx_verify.sh
3
4# 优化（默认 OPT=1，全栈优化）
5bash run_rocketx_verify.sh
6
7# 修改单卡 batch 为 16
8OPT=0 PER_GPU_BATCH=16 bash run_rocketx_verify.sh
9
10# 跑一小段用于性能测试（120 步）
11OPT=0 STEPS=120 NUM_EPOCHS=1 bash run_rocketx_verify.sh多机训练（2 节点示例）
两节点使用同一脚本，仅 NODE_RANK / MASTER_ADDR 不同（node0 先起，承载 rendezvous store）：
1# node0（rank0）
2MASTER_ADDR=<node0_ip> NODE_RANK=0 NNODES=2 GPUS_PER_NODE=8 OUTPUT_ROOT=<共享FS>/outputs bash run_rocketx_verify.sh
3
4# node1
5MASTER_ADDR=<node0_ip> NODE_RANK=1 NNODES=2 GPUS_PER_NODE=8 OUTPUT_ROOT=<共享FS>/outputs bash run_rocketx_verify.sh百舸平台会注入
MASTER_ADDR/MASTER_PORT/RANK（节点号）/WORLD_SIZE（节点数），脚本对这四个变量做兼容读取，百舸训练任务无需额外传参；显式拼写的NODE_RANK/NNODES优先于注入值。 容器要求：每个节点--network host、--ipc=host、--privileged、--shm-size足够大；若用--gpus "device=..."限定可见 GPU，必须同时给-e CUDA_VISIBLE_DEVICES=<容器内局部编号>。
多机正确性保障：采样器基于全局 rank / world_size 切分样本；只有 node 0 的 local rank 0 创建 work_dir、写 train.log、备份代码（避免多节点在共享 FS 上互相干扰）；RDMA 设备做硬前置校验（缺失直接报错，避免静默退化成 socket 拉低吞吐）；非法拓扑（NODE_RANK ≥ NNODES / GPUS_PER_NODE < 1 / 多机却用 127.0.0.1）直接报错退出。
性能参考
以下数据为 RTX Pro 6000（95.6 GiB/卡）实测，PER_GPU_BATCH=35，取训练稳态数据。
单机 8 卡（GBS = 8 × 35 = 280）
| 指标 | 基线 OPT=0 |
优化 OPT=1 |
提升 |
|---|---|---|---|
| 单轮次迭代速度（第 100 step） | 3.807 s/it | 2.261 s/it | 1.68×（−40.6%） |
| 单轮次迭代速度（稳态中位数，iter 6–120） | 3.920 s/it | 2.254 s/it | 1.74×（−42.5%） |
| 吞吐（样本级，按稳态中位数换算） | 71.4 samples/s | 124.2 samples/s | +73.9% |
峰值显存（nvidia-smi 实测，8 卡取最大） |
93.3 G / 95.6 G（97.6%） | 91.7 G / 95.6 G（95.9%） | −1.6 G |
data_load 中位数 |
20 ms | 6 ms | −70% |
| 单个权重大小 | 496.6 MB | 496.6 MB | — |
双机 16 卡（GBS = 2 × 8 × 35 = 560，RoCE v2 / RDMA）
| 指标 | 基线 OPT=0 |
优化 OPT=1 |
|---|---|---|
| 单步耗时（稳态中位） | 4.006 s/it | 2.295 s/it |
| 吞吐（样本级） | 139.8 samples/s | 244.0 samples/s |
峰值显存（nvidia-smi 实测，16 卡取最大） |
93.3 GiB | 91.2 GiB |
| 扩展比（单机单卡单步耗时 / 双机单卡单步耗时，阈值 0.9） | 0.98 | 0.98 |
扩展比 0.98 说明从单机 8 卡扩到双机 16 卡（GBS 随卡数翻倍），单卡单步耗时基本不变，横向扩展几乎无损。 以上数据为特定测试环境下的参考值。实际性能会受到资源规格、数据读取方式、网络环境、任务参数等因素影响。
loss 收敛对齐
优化前后 loss 收敛曲线不受影响（同一次 120-iter 对比）：
| 对比项 | 基线 OPT=0 |
优化 OPT=1 |
差异 |
|---|---|---|---|
| loss @ iter 120 | 41.0784 | 39.5046 | −3.8%（同一水平） |
| 收敛趋势 | 74.7 → 41.1 | 81.9 → 39.5 | 同趋势、同水平 |
iter 20 之后两者进入同一水平带（约 41），Δ 在 ±3% 内正负抖动，说明优化不改变收敛性。
断点续训
- 保存位置：
<OUTPUT_ROOT>/<model.type>/<RUN_NAME>/（默认outputs/SparseBEV/<RUN_NAME>/），文件命名iter_<iter>.pth，同时维护软链latest.pth - 保存频次：每
CKPT_SAVE_INTERVAL个 iteration 存一次，训练最后一步必定保存；CKPT_MAX_KEEP控制保留个数 - 单机续训：
1# 示例：每 1000 iter 存一次，保留 3 个
2CKPT_SAVE_INTERVAL=1000 CKPT_MAX_KEEP=3 bash run_rocketx_verify.sh
3
4# 从断点继续训练
5RESUME_FROM=outputs/SparseBEV/myrun/iter_1000.pth bash run_rocketx_verify.shRESUME_FROM 非空时脚本会校验文件存在，并令 work_dir = dirname(RESUME_FROM)（续训沿用原目录）；train.py 走 runner.resume()，恢复模型权重 + 完整 AdamW optimizer state + iteration 计数 + fp16 loss scaler + 学习率调度位置。多机续训时 OUTPUT_ROOT 必须是所有节点同路径可见的共享 FS。
常见问题
2026-11-05 前后版本说明
- 构建必须使用 legacy builder（
DOCKER_BUILDKIT=0 docker build）：RocketX 的setup.py用torch.cuda.is_available()决定是否编译 CUDA 扩展，BuildKit 的 RUN 沙箱看不到 GPU 会静默产出不含_rocketx_ops的 wheel（性能回退约 +7%）。镜像已内置构建期前置检查，拿不到 GPU / 算子缺失时直接 fail； run_rocketx_verify.sh中OPT=1分支默认ENABLE_ROCKETX_X=1（显式传入值优先），因此OPT=1 ENABLE_ROCKETX_TF32=0 bash run_rocketx_verify.sh可以正常只关 TF32。
日志里 mem_peak 和实际显存对不上？
训练日志里的 mem_peak 是 torch.cuda.max_memory_allocated()（只统计 rank0 卡上仍存活的 tensor 高水位），与 nvidia-smi 实测值差约 27 G，不要用日志 mem_peak 判断显存余量。显存吃紧时给基线加 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True 可显著降低碎片（实测全局峰值 93.3 G → 78.3 G，该变量需在 docker run 时注入）。
优化后 loss 与基线不一致，是不是有 bug？
前 20 iter 的逐 iter 数值差异（如 iter 1：74.67 → 81.85）来自 ENABLE_ROCKETX_PIN_MEMORY：pin_memory 补丁探测 collate 返回类型时会真实消费一个 batch，且 DataLoader 初始化会从全局 torch RNG 取一次 base_seed，导致「同一分布、不同随机序列」——属于正常的 shuffle/增广随机性，不是数值偏差。如果严格逐位对照，设置 ENABLE_ROCKETX_PIN_MEMORY=0 即可（此时与基线逐位一致）。其他优化项：P6-a/P6-b（关闭激活重计算）与基线逐位一致；P5-b（MSMV 反向算子）近似中性（Δ≈0.004%）；TF32 与匈牙利匹配向量化对 iter 1 无影响；P0 DN 分桶使 iter 1 变化约 −0.38%（DN 是训练期辅助监督，推理时整组丢弃，不改变主损失与优化目标）。
多机启动超时/未正常开始
- 多机两端需正确设置
MASTER_ADDR/NODE_RANK/NNODES，MASTER_PORT两端一致；node0 先启动。
本版本未覆盖项
- 节点规模验证到 2 节点 × 8 卡，更大规模（≥3 节点）的横向扩展性未实测。
参考
- SparseBEV 论文（纯视觉 3D 目标检测）：arXiv:2308.09244
- SparseBEV 官方仓库：https://github.com/MCG-NJU/SparseBEV
- RocketX 为镜像内置的零侵入训练加速库，以 wheel 形式安装、通过环境变量开关生效，无需修改训练代码
评价此篇文章