Nvidia Jetson Nano 运行 TensorRT：边缘 AI 图像识别推理优化实战指南

Nvidia Jetson Nano 是一款低功耗的边缘 AI 计算设备，专为实时推理任务设计，如图像识别。TensorRT 是 NVIDIA 的开源深度学习推理优化框架，能显著提升模型在边缘设备上的运行速度和效率。本指南将逐步介绍如何在 Jetson Nano 上部署和优化 TensorRT 模型，实现高效的图像识别推理。内容基于真实实践，确保可靠。

1. 环境准备：安装和配置 TensorRT

在 Jetson Nano 上，TensorRT 通常预装于 JetPack SDK 中。确保系统为最新版本：

• 更新系统：通过 SSH 登录 Jetson Nano，运行命令：

  sudo apt update && sudo apt upgrade -y
  

• 检查 TensorRT 版本：运行 dpkg -l | grep tensorrt 确认安装（如 TensorRT 8.x）。
• 安装依赖：安装 Python 库：

  pip install numpy pillow pycuda
  

• 验证安装：运行 Python 脚本导入 import tensorrt as trt，无错误即成功。

2. 模型准备：转换为 TensorRT 引擎

TensorRT 需要将训练好的模型（如 PyTorch 或 TensorFlow 模型）转换为优化引擎。推荐使用 ONNX 格式作为中介：

• 步骤：
  1. 导出模型为 ONNX：在训练环境中，使用框架工具（如 PyTorch 的 torch.onnx.export）。
  2. 在 Jetson Nano 上，使用 TensorRT 的 ONNX 解析器转换：

     import tensorrt as trt
     
     # 创建 TensorRT 构建器
     TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
     builder = trt.Builder(TRT_LOGGER)
     network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
     parser = trt.OnnxParser(network, TRT_LOGGER)
     
     # 加载 ONNX 模型
     with open("model.onnx", "rb") as f:
         parser.parse(f.read())
     
     # 配置优化选项（如 FP16 精度）
     config = builder.create_builder_config()
     config.set_flag(trt.BuilderFlag.FP16)  # 启用半精度浮点
     
     # 构建引擎并保存
     engine = builder.build_engine(network, config)
     with open("model.trt", "wb") as f:
         f.write(engine.serialize())
     

• 优化原理：TensorRT 通过层融合和权重剪枝减少计算量。例如，卷积层优化可表示为： $$ y = f(W \ast X + b) $$ 其中 $W$ 为权重，$X$ 为输入，$b$ 为偏置，$f$ 为激活函数。TensorRT 自动简化此过程。

3. 推理优化实战：图像识别代码示例

以下 Python 代码演示如何在 Jetson Nano 上加载 TensorRT 引擎，进行图像识别（以 ResNet 模型为例）。优化重点包括预处理、推理执行和后处理：

• 完整代码：

  import tensorrt as trt
  import pycuda.driver as cuda
  import pycuda.autoinit
  import numpy as np
  from PIL import Image
  
  # 加载 TensorRT 引擎
  def load_engine(engine_path):
      with open(engine_path, "rb") as f, trt.Runtime(trt.Logger(trt.Logger.WARNING)) as runtime:
          return runtime.deserialize_cuda_engine(f.read())
  
  # 预处理图像：调整大小、归一化
  def preprocess_image(image_path):
      img = Image.open(image_path).resize((224, 224))  # 输入尺寸匹配模型
      img = np.array(img).astype(np.float32) / 255.0  # 归一化到 [0,1]
      img = np.transpose(img, (2, 0, 1))  # HWC 到 CHW
      return np.expand_dims(img, axis=0)  # 添加批处理维度
  
  # 运行推理
  def infer(engine, input_data):
      # 创建执行上下文
      context = engine.create_execution_context()
      # 分配 GPU 内存
      input_shape = engine.get_binding_shape(0)
      output_shape = engine.get_binding_shape(1)
      input_mem = cuda.mem_alloc(input_data.nbytes)
      output_mem = cuda.mem_alloc(np.prod(output_shape) * 4)  # 假设输出为 float32
      # 数据拷贝到 GPU
      cuda.memcpy_htod(input_mem, input_data)
      # 执行推理
      context.execute_v2(bindings=[int(input_mem), int(output_mem)])
      # 拷贝结果回 CPU
      output_data = np.empty(output_shape, dtype=np.float32)
      cuda.memcpy_dtoh(output_data, output_mem)
      return output_data
  
  # 主函数：图像识别
  if __name__ == "__main__":
      engine_path = "model.trt"  # 替换为你的引擎文件
      image_path = "test.jpg"    # 测试图像
      engine = load_engine(engine_path)
      input_data = preprocess_image(image_path)
      output = infer(engine, input_data)
      predicted_class = np.argmax(output)  # 假设输出为分类概率
      print(f"预测类别: {predicted_class}")
  

• 优化说明：
  • 预处理加速：使用 PIL 和 NumPy 在 CPU 上快速处理图像，减少 GPU 瓶颈。
  • 内存管理：通过 PyCUDA 直接操作 GPU 内存，避免不必要的拷贝。
  • 精度设置：在引擎构建时启用 FP16（半精度），可减少内存占用和计算时间，推理速度提升可达 2-3 倍。

4. 关键优化技巧

在边缘设备上，推理性能受限于算力和内存。以下是实战优化策略：

• 精度调整：使用 FP16 或 INT8 量化。INT8 进一步减少模型大小和延迟，但需校准：
  • 数学基础：量化公式为 $x_{\text{int}} = \text{round}(x / \text{scale})$，其中 $\text{scale}$ 为缩放因子。
  • 在 TensorRT 配置中添加：

    config.set_flag(trt.BuilderFlag.INT8)
    config.int8_calibrator = ...  # 使用校准数据集
    

• 批处理优化：增加批处理大小（batch size）以并行处理多个图像，提高吞吐量。但需权衡内存：
  • 在引擎构建时设置最大批处理尺寸：builder.max_batch_size = 4。
• 层融合：TensorRT 自动融合卷积和激活层，减少 kernel 调用次数。例如，融合 ReLU 层： $$ y = \max(0, W \ast X + b) $$
• 实时监控：使用 Jetson Nano 的 tegrastats 工具监控 GPU 和 CPU 使用率，调整参数：

  tegrastats --interval 1000  # 每秒输出资源使用
  

• 预期性能：在 Jetson Nano 上，优化后推理延迟可降至 10-50ms（取决于模型复杂度），相比原生框架提升 5-10 倍。

5. 性能评估和调试

• 基准测试：使用 Python 的 time 模块测量推理延迟：

  import time
  start = time.time()
  output = infer(engine, input_data)
  latency = (time.time() - start) * 1000  # 毫秒
  print(f"推理延迟: {latency:.2f} ms")
  

• 常见问题：
  • 内存不足：减少批处理大小或使用 INT8 量化。
  • 精度损失：在 FP16 模式下，添加正则化层或使用混合精度训练。
• 工具推荐：使用 NVIDIA Nsight Systems 分析性能瓶颈。

结论

通过本指南，您可以在 Nvidia Jetson Nano 上高效部署 TensorRT 模型，实现图像识别的边缘推理优化。实战中，FP16 或 INT8 量化是关键，能显著提升速度和能效。优化后，Jetson Nano 可处理实时视频流（例如 30 FPS），适用于智能监控或 IoT 应用。建议从简单模型（如 MobileNet）开始测试，逐步扩展到复杂场景。如有问题，可参考 TensorRT 官方文档或社区资源。