1. 项目背景与挑战
在自动驾驶领域,BEVFormer作为一种基于纯视觉的鸟瞰图感知框架,近年来在3D目标检测和语义地图分割任务中表现出色。然而,将其部署到边缘计算设备如Jetson Orin NX上时,面临着诸多技术挑战。
我最近在Jetson Orin NX(JetPack 6.2)上成功部署了BEVFormer模型,整个过程可谓一波三折。主要难点在于:
- 官方BEVFormer代码基于较旧的PyTorch 1.9.1和CUDA 11.1开发
- Jetson Orin NX预装的是CUDA 12.6和PyTorch 2.5
- 硬件架构差异(x86-64 vs ARM aarch64)导致预编译轮子不可用
2. 环境准备与方案选择
2.1 硬件与基础环境配置
我的开发环境分为两部分:
PC端配置:
- GPU: RTX 5080 (SM_120架构)
- CUDA: 12.9
- PyTorch: 2.7.0
- Python: 3.10
- OS: Ubuntu 22.04 x86-64
Jetson Orin NX配置:
- JetPack: 6.2
- CUDA: 12.6
- TensorRT: 10.3.0.30
- cuDNN: 9.3.0
- Python: 3.10
- OS: Ubuntu 22.04 aarch64
2.2 部署方案评估
面对版本不兼容问题,我评估了两种方案:
方案一:降级JetPack
- 安装JetPack 5.1以匹配官方BEVFormer要求
- 问题:JetPack 5.1不支持Ubuntu 22.04
- 结论:放弃此方案
方案二:适配高版本环境
- 保持JetPack 6.2
- 从源码编译PyTorch 2.7.0
- 使用修改版的BEVFormer代码
- 最终选择此方案
3. PyTorch 2.7.0的ARM编译
3.1 准备工作
由于官方未提供ARM架构的PyTorch 2.7.0预编译包,必须从源码编译:
bash复制git clone https://github.com/pytorch/pytorch
cd pytorch
git submodule sync
git submodule update --init --recursive
3.2 编译配置
创建setup.sh配置文件,关键参数如下:
bash复制unset NVCC_APPEND_FLAGS CMAKE_CUDA_FLAGS
export USE_CUDA=1
export USE_CUDNN=1
export USE_NCCL=0
export USE_DISTRIBUTED=ON # 必须为ON才能支持分布式训练
export USE_TENSORRT=1
export TORCH_CUDA_ARCH_LIST="8.7" # Orin NX的Ampere架构
export MAX_JOBS=10 # 防止OOM
export PYTORCH_BUILD_VERSION=2.7.0
export PYTORCH_BUILD_NUMBER=0
export CC=/usr/bin/gcc-9
export CXX=/usr/bin/g++-9
3.3 关键编译选项解析
| 选项 | 设置 | 原因 |
|---|---|---|
| USE_DISTRIBUTED | ON | BEVFormer测试脚本需要分布式支持 |
| USE_NCCL | OFF | Jetson平台不支持NCCL |
| TORCH_CUDA_ARCH_LIST | 8.7 | 匹配Orin NX的Ampere架构 |
| NVCC_APPEND_FLAGS | 特殊设置 | 解决ARM架构下的编译问题 |
3.4 编译与安装
执行编译命令:
bash复制source setup.sh
python setup.py bdist_wheel
编译完成后,安装生成的wheel包:
bash复制cd dist && pip install torch-2.7.0-cp310-cp310-linux_aarch64.whl
注意:完整编译可能需要4-6小时,建议使用screen或tmux保持会话
4. 配套库的编译安装
4.1 torchvision编译
bash复制git clone https://github.com/pytorch/vision.git
export FORCE_CUDA=1
export TORCH_CUDA_ARCH_LIST="8.7"
python setup.py install
4.2 Detectron2安装
由于网络问题,采用本地安装:
bash复制git clone https://github.com/facebookresearch/detectron2.git
pip install -e . --no-build-isolation
4.3 MMCV定制编译
使用修改版的MMCV 1.7.2:
bash复制cd mmcv
MMCV_WITH_OPS=1 pip install . --no-cache-dir -v --no-build-isolation
4.4 MMDetection3D安装
安装修改版的MMDetection3D 1.0.0rc6:
bash复制pip install -e .
5. 代码适配与问题解决
5.1 分布式训练适配
在tools/test.py中修改分布式相关代码:
python复制if not distributed:
model = MMDataParallel(model, device_ids=[0])
outputs = single_gpu_test(model, data_loader, args.show, args.show_dir)
else:
from torch.nn.parallel import DistributedDataParallel
model = DistributedDataParallel(
model.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False)
outputs = custom_multi_gpu_test(model, data_loader, args.tmpdir, args.gpu_collect)
5.2 常见编译问题解决
-
CUDA架构不匹配:
- 错误:
Unsupported gpu architecture 'compute_xx' - 解决:确保
TORCH_CUDA_ARCH_LIST设置为"8.7"
- 错误:
-
内存不足:
- 错误:
g++: fatal error: Killed signal terminated program cc1plus - 解决:减少
MAX_JOBS数量,清理内存
- 错误:
-
Python头文件缺失:
- 错误:
Python.h not found - 解决:安装
python3-dev包
- 错误:
6. 部署验证与性能测试
6.1 功能验证
运行BEVFormer测试脚本:
bash复制python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_FILE} --eval bbox
验证输出:
- 确认生成测试结果文件
- 检查评估指标是否符合预期
6.2 性能优化建议
-
TensorRT加速:
- 将模型转换为TensorRT引擎
- 利用Orin NX的DLA加速
-
量化部署:
- 采用FP16或INT8量化
- 使用PyTorch的量化工具
-
内存优化:
- 调整batch size
- 启用CUDA graph
7. 经验总结与资源分享
这次部署过程中积累了几个关键经验:
-
版本兼容性矩阵至关重要,建议提前规划好整个软件栈的版本
-
ARM架构下的源码编译非常耗时,建议:
- 使用ccache加速后续编译
- 准备编译好的wheel包供团队共享
-
我已将编译好的torch 2.7.0 aarch64 wheel包分享在网盘:
- 链接: https://pan.baidu.com/s/1tO-l_RIVLb3cXS9H0HmOVA
- 提取码: u249
-
对于后续的模型量化,建议关注:
- TensorRT的FP16/INT8量化
- PyTorch自带的量化工具
- 量化后的精度验证策略
