将训练好的 YOLOv8 模型轻量化并导出为 ONNX/TensorRT 格式,最终部署到 Jetson Nano 这类嵌入式设备上加速推理,这是实际落地中非常关键的一步,能显著提升模型在边缘设备上的运行速度。
下面我会从模型导出、TensorRT 优化、Jetson Nano 部署三个核心环节,给出完整的可执行步骤和代码,确保你能一步步落地。
一、核心原理
- ONNX:是一种跨框架的模型格式,作为中间桥梁,能实现 PyTorch/TensorFlow 等框架到 TensorRT 的转换。
- TensorRT:NVIDIA 推出的高性能推理引擎,通过层融合、精度校准(FP16/INT8)、内核自动调优等方式优化模型,大幅提升 NVIDIA 设备(如 Jetson Nano)的推理速度。
- 轻量化核心:除了格式转换,还包括精度降低(从 FP32 到 FP16/INT8)、模型裁剪等,在精度损失可控的前提下提升速度。
二、完整实现步骤
前置条件
- Jetson Nano 环境准备:
- 刷入 JetPack 系统(推荐 JetPack 5.0+,自带 TensorRT、CUDA)。
- 确认已安装
ultralytics、onnx、tensorrt:bash
运行
# Jetson Nano 安装 ultralytics pip3 install ultralytics onnx # TensorRT 随 JetPack 自带,无需额外安装,验证是否安装成功 python3 -c "import tensorrt; print(tensorrt.__version__)"
- 准备好训练好的 YOLOv8 模型(
best.pt)。
步骤 1:导出 YOLOv8 模型为 ONNX 格式
YOLOv8 内置了一键导出功能,无需手动编写转换代码,直接导出 ONNX 格式(兼容 TensorRT)。
python
运行
from ultralytics import YOLO # 加载训练好的 YOLOv8 模型 model = YOLO("best.pt") # 导出为 ONNX 格式(关键参数说明) # format: 导出格式,指定为 onnx # imgsz: 输入图像尺寸,需和训练时一致(如 640) # opset: ONNX 算子版本,推荐 12+(兼容 TensorRT) # dynamic: 动态输入尺寸(可选,嵌入式设备建议固定尺寸以提升速度) # half: 导出 FP16 精度(轻量化核心,减少显存占用) model.export( format="onnx", imgsz=640, opset=12, dynamic=False, half=True # 导出 FP16 精度,Jetson Nano 支持 )执行后会在模型同目录生成best.onnx文件,这是中间格式模型。
步骤 2:将 ONNX 转为 TensorRT 引擎(核心加速)
TensorRT 引擎(.engine文件)是针对 NVIDIA 硬件优化后的模型,推理速度最快。有两种方式实现:
方式 1:使用 YOLOv8 内置的 TensorRT 导出(最简单)
python
运行
# 直接从 YOLOv8 模型导出 TensorRT 引擎 model.export( format="engine", imgsz=640, device=0, # Jetson Nano 的 GPU 设备号 half=True # FP16 精度 )执行后生成best.engine文件,直接用于推理。
方式 2:手动用 TensorRT API 转换(更灵活,适合自定义优化)
python
运行
import tensorrt as trt import os # TensorRT 日志设置 TRT_LOGGER = trt.Logger(trt.Logger.INFO) EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) def build_engine(onnx_file_path, engine_file_path): # 构建 TensorRT 引擎 builder = trt.Builder(TRT_LOGGER) network = builder.create_network(EXPLICIT_BATCH) parser = trt.OnnxParser(network, TRT_LOGGER) config = builder.create_builder_config() # 设置最大工作空间(Jetson Nano 建议 1GB 以内) config.max_workspace_size = 1 << 30 # 1GB # 启用 FP16 精度(Jetson Nano 支持) config.set_flag(trt.BuilderFlag.FP16) # 解析 ONNX 文件 with open(onnx_file_path, "rb") as f: if not parser.parse(f.read()): for error in range(parser.num_errors): print(parser.get_error(error)) return None # 构建并保存引擎 serialized_engine = builder.build_serialized_network(network, config) with open(engine_file_path, "wb") as f: f.write(serialized_engine) return engine_file_path # 转换 ONNX 到 TensorRT 引擎 onnx_file = "best.onnx" engine_file = "best.engine" if not os.path.exists(engine_file): build_engine(onnx_file, engine_file) print(f"TensorRT 引擎已生成:{engine_file}") else: print(f"引擎文件已存在:{engine_file}")步骤 3:在 Jetson Nano 上用 TensorRT 引擎推理
使用生成的.engine文件进行推理,速度远快于原始.pt模型:
python
运行
from ultralytics import YOLO # 加载 TensorRT 引擎模型 model = YOLO("best.engine") # 推理单张图像(Jetson Nano 上实测速度比 PT 模型快 2-5 倍) results = model("test.jpg", imgsz=640) # 解析位置结果(和之前 PT 模型推理代码完全兼容) for r in results: boxes = r.boxes for box in boxes: x1, y1, x2, y2 = box.xyxy[0].tolist() cls_name = model.names[box.cls[0].item()] conf = box.conf[0].item() print(f"类别:{cls_name}, 置信度:{conf:.2f}, 位置:({x1:.1f},{y1:.1f})-({x2:.1f},{y2:.1f})") # 可视化结果 results[0].save("result_trt.jpg")步骤 4:进阶优化(INT8 精度,极致轻量化)
如果追求更快速度,可将模型转为 INT8 精度(需准备校准数据集):
python
运行
# YOLOv8 导出 INT8 精度的 TensorRT 引擎 model.export( format="engine", imgsz=640, device=0, int8=True, # 启用 INT8 精度 data="data.yaml" # 校准数据集配置文件,用于 INT8 量化 )注意:INT8 精度会有轻微精度损失,但速度比 FP16 再提升 30%-50%,适合对精度要求不极致的场景。
三、Jetson Nano 部署注意事项
- 显存限制:Jetson Nano 只有 4GB 显存,建议使用 YOLOv8n(轻量版)、imgsz=640 或更小(如 480)。
- 电源供电:必须用 5V/3A 电源,否则会因供电不足导致推理卡顿或设备重启。
- 散热:Jetson Nano 推理时发热明显,需加装散热风扇,避免因过热降频。
- 推理速度参考:
- YOLOv8n + FP32(PT 模型):~5 FPS
- YOLOv8n + FP16(TensorRT):~15 FPS
- YOLOv8n + INT8(TensorRT):~20 FPS
总结
- 模型轻量化到 Jetson Nano 部署的核心流程:
YOLOv8 PT 模型 → ONNX 中间格式 → TensorRT 引擎,其中 TensorRT 是嵌入式设备加速的关键。 - 优先选择FP16 精度导出 TensorRT 引擎,在精度损失可控的前提下,能让推理速度提升 2-5 倍。
- Jetson Nano 部署需注意显存、供电和散热,建议使用轻量版模型(YOLOv8n)和固定输入尺寸。