Jetson板卡PyTorch推理环境配置全攻略

一、环境配置前的硬件与系统准备

1.1 Jetson板卡选型与性能对比

Jetson系列板卡包含Nano（4GB/2GB）、TX2、Xavier NX及AGX Xavier等型号，其核心差异在于GPU算力（如Nano为0.5TFLOPS，AGX Xavier达32TFLOPS）和内存带宽。对于PyTorch推理任务，建议选择至少4GB内存的型号（如Nano 4GB或Xavier系列），以避免模型加载时的内存溢出问题。

1.2 系统镜像安装与初始化

NVIDIA官方提供JetPack SDK作为Jetson设备的系统镜像，需通过以下步骤完成安装：

1. 从NVIDIA官网下载对应板卡的JetPack镜像（如L4T R32.x或R35.x）
2. 使用balenaEtcher等工具将镜像写入SD卡（Nano系列）或通过NVIDIA SDK Manager进行本地安装
3. 首次启动后，通过sudo apt update && sudo apt upgrade更新系统包

1.3 CUDA与cuDNN版本匹配

Jetson板卡内置的CUDA工具包版本需与PyTorch兼容。例如：

• JetPack 4.6（L4T R32.6.1）默认集成CUDA 10.2和cuDNN 8.0
• JetPack 5.x（L4T R35.x）升级至CUDA 11.4和cuDNN 8.2
  可通过nvcc --version和cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR命令验证版本。

二、PyTorch框架安装与验证

2.1 官方预编译包安装（推荐）

NVIDIA为Jetson平台提供了预编译的PyTorch轮子（wheel），安装步骤如下：

# 以JetPack 4.6为例
wget https://nvidia.box.com/shared/static/fjtbno0v9i67ddu0jlpggjkhq07tldru.whl -O torch-1.8.0-cp36-cp36m-linux_aarch64.whl
pip3 install torch-1.8.0-cp36-cp36m-linux_aarch64.whl
pip3 install torchvision==0.9.0  # 需匹配PyTorch版本

关键点：需严格对应Python版本（如CP36对应Python 3.6）、JetPack版本和架构（aarch64）。

2.2 源码编译安装（高级用户）

若需自定义PyTorch版本，可通过源码编译实现，但过程复杂且耗时较长：

1. 安装依赖：sudo apt install git cmake ninja-build libopenblas-dev libjpeg-dev
2. 克隆PyTorch源码：git clone --recursive https://github.com/pytorch/pytorch
3. 切换至指定版本标签（如v1.8.0）
4. 修改setup.py中的CUDA路径（如/usr/local/cuda-10.2）
5. 执行编译：USE_CUDA=1 USE_NNPACK=0 python3 setup.py install

2.3 环境验证

安装完成后，通过以下Python代码验证环境：

import torch
print(torch.__version__)  # 应输出安装版本（如1.8.0）
print(torch.cuda.is_available())  # 应输出True
x = torch.rand(3, 3).cuda()
print(x.device)  # 应输出cuda:0

三、推理优化技巧

3.1 TensorRT加速集成

Jetson板卡支持通过TensorRT优化PyTorch模型：

1. 导出ONNX模型：

   dummy_input = torch.randn(1, 3, 224, 224)
   torch.onnx.export(model, dummy_input, "model.onnx")

2. 使用trtexec工具转换为TensorRT引擎：

   trtexec --onnx=model.onnx --saveEngine=model.trt --fp16

3. 在PyTorch中加载TensorRT引擎（需通过自定义算子或ONNX Runtime）。

3.2 内存与性能调优

• 半精度推理：启用model.half()和torch.cuda.FloatTensor.half()减少内存占用
• 批处理优化：通过torch.nn.DataParallel实现多GPU并行（Xavier系列支持双GPU）
• 动态输入尺寸：使用torch.jit.trace固定输入尺寸以避免动态形状开销

3.3 常见问题解决

1. CUDA内存不足：

   • 减小批处理大小（batch size）
   • 启用torch.backends.cudnn.benchmark = True
   • 使用nvidia-smi监控GPU内存占用

2. PyTorch与TensorRT版本冲突：

   • 确保TensorRT版本与JetPack匹配（如JetPack 4.6对应TensorRT 7.1）
   • 避免混合安装不同来源的PyTorch和TensorRT包

3. 模型加载失败：

   • 检查模型架构与PyTorch版本兼容性
   • 使用torch.load(..., map_location='cuda:0')强制指定设备

四、完整工作流示例

以下是一个从模型加载到推理的完整示例：

import torch
import torchvision.models as models
# 1. 加载预训练模型
model = models.resnet18(pretrained=True).eval().half().cuda()
# 2. 准备输入数据
input_tensor = torch.randn(1, 3, 224, 224).half().cuda()
# 3. 执行推理
with torch.no_grad():
    output = model(input_tensor)
# 4. 输出结果
print("Output shape:", output.shape)
print("Top-5 classes:", torch.topk(output, 5).indices.cpu().numpy())

五、进阶资源推荐

1. NVIDIA官方文档：

   • Jetson Developer Forum
   • JetPack SDK文档

2. 开源项目参考：

   • PyTorch for Jetson（维护预编译包）
   • TensorRT示例

3. 性能分析工具：

   • nvprof：CUDA内核级性能分析
   • torch.autograd.profiler：PyTorch操作级分析

通过以上步骤，开发者可在Jetson板卡上高效部署PyTorch推理任务，平衡性能与资源占用。实际项目中，建议结合具体模型复杂度和硬件规格进行针对性优化。