YOLOv5模型在Jetson Nano上的TensorRT部署

YOLOv5模型在Jetson Nano上的TensorRT部署

边缘智能的落地挑战：从训练到推理的鸿沟

在嵌入式AI设备日益普及的今天，一个常见但棘手的问题浮出水面：我们能在PC上轻松训练出高精度的目标检测模型，却常常卡在“如何让它真正在小设备上跑起来”这一步。比如Jetson Nano这样的边缘计算平台，虽然集成了GPU，但算力和内存资源极为有限，直接运行PyTorch模型几乎不可能。

YOLOv5作为目标检测领域的明星模型，以其简洁架构和高效推理著称。然而，它的.pt文件本质上是为训练设计的，包含大量冗余操作和动态图结构，并不适合部署。要让YOLOv5在Jetson Nano上实现稳定实时推理，必须借助NVIDIA的TensorRT——这才是打通“实验室”与“现场应用”的关键一环。

为什么是 Jetson Nano + TensorRT？

Jetson Nano虽小，五脏俱全。它搭载了128核Maxwell GPU和四核ARM Cortex-A57 CPU，支持CUDA、cuDNN和TensorRT，功耗仅5W左右，非常适合用于机器人视觉、智能监控或工业质检等场景。但其4GB内存（开发版）也意味着任何部署方案都必须高度优化。

而TensorRT正是为此而生。它不仅能将FP32模型转换为FP16甚至INT8以压缩体积、提升速度，还能执行层融合（layer fusion）、内核自动调优（kernel autotuning）等底层优化。更重要的是，它生成的是静态计算图，避免了PyTorch运行时的开销，在资源受限设备上性能提升可达数倍。

举个例子：原生PyTorch下的YOLOv5s在Nano上可能只能勉强跑到2~3 FPS，而经过TensorRT优化后，同一模型在640×640输入下可达接近10 FPS，帧率提升超过300%，且CPU占用更低，发热更可控。

部署路径重构：不再依赖完整框架

许多初学者尝试直接在Jetson Nano上安装Ultralytics库并加载.pt模型进行推理，结果往往是内存溢出或卡顿严重。这不是代码问题，而是思路偏差——我们不该把训练环境搬到边缘端，而应只保留“推理”所需的最小闭环。

理想路径如下：

YOLOv5 .pt 模型 → 提取权重生成 .wts 文件 → 在目标设备使用 tensorrtx 编译为 .engine 引擎 → 用 C++ 或轻量Python脚本调用

其中核心工具链来自开源项目wang-xinyu/tensorrtx，它提供了一套清晰的C++模板，专为YOLO系列适配TensorRT而设计。相比官方export format=engine命令，这种方式更灵活、兼容性更强，尤其适合跨平台部署。

✅ 注：本文基于Ultralytics官方YOLOv5 v6.1版本验证，适用于主流自定义数据集迁移。

硬件准备：别让供电毁了整个项目

Jetson Nano看似即插即用，实则对硬件要求极为敏感。以下是成功部署的前提条件：

• 推荐使用4GB版本开发套件
  2GB版本显存紧张，运行yolov5m/l/x极易OOM。
• MicroSD卡 ≥ 64GB，Class 10及以上
  系统+依赖+模型文件合计可能超30GB。
• 必须配备5V/4A直流电源
  使用USB供电时易因电流不足导致系统重启或降频。
• 主动散热不可少
  加装风扇+金属散热片，否则持续负载下GPU会因过热降频至50%以下。
• 外接HDMI显示器、键盘鼠标一套
  首次配置无法免驱，远程连接需后续手动开启。
• 网络接入方式任选其一：网线 or 支持Linux驱动的USB无线网卡。

⚠️ 特别提醒：若使用桶形接口供电，请确保J48跳线帽已正确安装（靠近GPIO排针处），否则无法启用外部电源模式。

软件环境搭建：精准匹配版本至关重要

Jetson Nano不支持通用Ubuntu镜像，必须刷写NVIDIA官方定制的L4T系统。推荐组合如下：

前往 NVIDIA官网 下载完整SDK，使用SD Card Image Writer写入TF卡后启动设备。

首次开机完成基础设置后，建议执行以下优化：

# 更换国内源加速apt更新（清华源） sudo sed -i 's/ports.ubuntu.com/mirrors.tuna.tsinghua.edu.cn/g' /etc/apt/sources.list sudo apt update && sudo apt upgrade -y # 安装必要工具链 sudo apt install vim git wget build-essential cmake libopencv-dev python3-pip -y # 关键依赖：pycuda（用于GPU内存交互） sudo pip3 install pycuda # 验证环境是否正常 nvcc --version # 应输出 CUDA 10.2 dpkg -l | grep tensorrt # 查看 TensorRT 是否安装

第一步：提取模型权重（.wts）

尽管新版Ultralytics支持export format=engine，但在Jetson Nano上直接生成往往失败，原因包括：
- 架构差异（x86 vs ARM）
- 插件算子未注册
- 动态shape支持不一致

更稳妥的方式是在PC端导出纯权重文件.wts，再传输至Nano侧编译。

在主机端操作

git clone https://github.com/ultralytics/yolov5.git cd yolov5 pip install -r requirements.txt # 下载 tensorrtx 提供的权重提取脚本 wget https://raw.githubusercontent.com/wang-xinyu/tensorrtx/master/yolov5/gen_wts.py # 以 yolov5s.pt 为例生成 .wts 文件 python gen_wts.py -w yolov5s.pt

成功后会生成yolov5s.wts，其格式为每行列出层名及对应权重值，不含任何计算逻辑，便于跨平台移植。

通过SCP等方式将其复制到Jetson Nano：

scp yolov5s.wts jetson@<nano-ip>:/home/jetson/tensorrtx/yolov5/

第二步：在 Jetson Nano 上构建 TensorRT 引擎

获取 tensorrtx 工程

git clone https://github.com/wang-xinyu/tensorrtx.git cd tensorrtx/yolov5

当前版本主要支持YOLOv5-v6.1及以下。若使用更新版模型，需手动调整网络结构定义。

放置权重并修改配置

将.wts文件放入当前目录，并根据实际任务调整参数。

修改类别数量

编辑yololayer.h中的常量：

static constexpr int CLASS_NUM = 80; // 默认COCO

替换为你训练时的类别数，例如四分类：

static constexpr int CLASS_NUM = 4;

调整输入分辨率（可选）

默认为640×640，可在yolov5.cpp中修改：

const int INPUT_H = 416; const int INPUT_W = 416;

⚠️ 输入尺寸应为32的倍数。减小分辨率可显著提速，但会影响小物体检测能力。

编译并序列化引擎

mkdir build && cd build cmake .. make -j4

若编译报错缺少头文件，请确认是否已安装OpenCV开发包（libopencv-dev）。

生成可执行文件后，执行序列化命令：

sudo ./yolov5 -s ../yolov5s.wts yolov5s.engine s

参数说明：
--s：表示序列化（serialize）
-../yolov5s.wts：输入权重路径
-yolov5s.engine：输出引擎名称
-s：模型类型（支持 s/m/l/x/s6/m6/l6/x6）

✅ 若无错误提示，则yolov5s.engine已生成，可用于后续推理调用。

第三步：Python调用TRT引擎实现轻量推理

虽然C++性能最优，但Python更适合快速验证和集成应用。下面是一个无需PyTorch的轻量级推理脚本。

创建yolov5_trt.py

import os import cv2 import time import numpy as np import pycuda.autoinit import pycuda.driver as cuda import tensorrt as trt import ctypes import random

设置全局变量

INPUT_W = 640 INPUT_H = 640 CONF_THRESH = 0.25 IOU_THRESHOLD = 0.45 # 替换为你的类别标签（顺序必须与训练一致） categories = ['person', 'bicycle', 'car', 'motorcycle']

绘图函数

def plot_one_box(x, img, color=None, label=None, line_thickness=None): tl = line_thickness or max(int((img.shape[0] + img.shape[1]) / 600), 1) color = color or [random.randint(0, 255) for _ in range(3)] c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3])) cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA) if label: tf = max(tl - 1, 1) t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)

核心类：YoLov5TRT

class YoLov5TRT: def __init__(self, engine_file_path): self.ctx = cuda.Device(0).make_context() stream = cuda.Stream() TRT_LOGGER = trt.Logger(trt.Logger.INFO) runtime = trt.Runtime(TRT_LOGGER) with open(engine_file_path, "rb") as f: engine = runtime.deserialize_cuda_engine(f.read()) context = engine.create_execution_context() host_inputs = [] cuda_inputs = [] host_outputs = [] cuda_outputs = [] bindings = [] for binding in engine: size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size dtype = trt.nptype(engine.get_binding_dtype(binding)) host_mem = cuda.pagelocked_empty(size, dtype) cuda_mem = cuda.mem_alloc(host_mem.nbytes) bindings.append(int(cuda_mem)) if engine.binding_is_input(binding): host_inputs.append(host_mem) cuda_inputs.append(cuda_mem) else: host_outputs.append(host_mem) cuda_outputs.append(cuda_mem) self.stream = stream self.context = context self.engine = engine self.host_inputs = host_inputs self.cuda_inputs = cuda_inputs self.host_outputs = host_outputs self.cuda_outputs = cuda_outputs self.bindings = bindings def infer(self, image): self.ctx.push() stream = self.stream context = self.context host_inputs = self.host_inputs cuda_inputs = self.cuda_inputs host_outputs = self.host_outputs cuda_outputs = self.cuda_outputs bindings = self.bindings input_image, origin_h, origin_w = self.preprocess(image) np.copyto(host_inputs[0], input_image.ravel()) cuda.memcpy_htod_async(cuda_inputs[0], host_inputs[0], stream) context.execute_async_v2(bindings=bindings, stream_handle=stream.handle) cuda.memcpy_dtoh_async(host_outputs[0], cuda_outputs[0], stream) stream.synchronize() self.ctx.pop() output = host_outputs[0] boxes, scores, class_ids = self.post_process(output, origin_h, origin_w) return boxes, scores, class_ids def preprocess(self, image): h, w, c = image.shape image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) r_w = INPUT_W / w r_h = INPUT_H / h if r_h > r_w: tw = INPUT_W th = int(r_w * h) tx1 = tx2 = 0 ty1 = int((INPUT_H - th) / 2) ty2 = INPUT_H - th - ty1 else: tw = int(r_h * w) th = INPUT_H tx1 = int((INPUT_W - tw) / 2) tx2 = INPUT_W - tw - tx1 ty1 = ty2 = 0 image_resized = cv2.resize(image_rgb, (tw, th)) image_padded = cv2.copyMakeBorder(image_resized, ty1, ty2, tx1, tx2, cv2.BORDER_CONSTANT, value=(128, 128, 128)) image_normalized = image_padded.astype(np.float32) / 255.0 image_chw = np.transpose(image_normalized, [2, 0, 1]) image_nchw = np.expand_dims(image_chw, axis=0) image_contiguous = np.ascontiguousarray(image_nchw) return image_contiguous, h, w def post_process(self, output, origin_h, origin_w): num_boxes = int(output[0]) pred = np.reshape(output[1:], (-1, 6))[:num_boxes] boxes_xywh = pred[:, :4] scores = pred[:, 4] class_ids = pred[:, 5] conf_mask = scores > CONF_THRESH boxes_xywh = boxes_xywh[conf_mask] scores = scores[conf_mask] class_ids = class_ids[conf_mask] boxes_xyxy = self.xywh2xyxy(origin_h, origin_w, boxes_xywh) keep = self.nms(boxes_xyxy, scores) return boxes_xyxy[keep], scores[keep], class_ids[keep] def xywh2xyxy(self, origin_h, origin_w, x): y = np.zeros_like(x) r_w = INPUT_W / origin_w r_h = INPUT_H / origin_h if r_h > r_w: y[:, 0] = x[:, 0] - x[:, 2] / 2 y[:, 2] = x[:, 0] + x[:, 2] / 2 y[:, 1] = x[:, 1] - x[:, 3] / 2 - (INPUT_H - r_w * origin_h) / 2 y[:, 3] = x[:, 1] + x[:, 3] / 2 - (INPUT_H - r_w * origin_h) / 2 y /= r_w else: y[:, 0] = x[:, 0] - x[:, 2] / 2 - (INPUT_W - r_h * origin_w) / 2 y[:, 2] = x[:, 0] + x[:, 2] / 2 - (INPUT_W - r_h * origin_w) / 2 y[:, 1] = x[:, 1] - x[:, 3] / 2 y[:, 3] = x[:, 1] + x[:, 3] / 2 y /= r_h return y def nms(self, boxes, scores): x1 = boxes[:, 0] y1 = boxes[:, 1] x2 = boxes[:, 2] y2 = boxes[:, 3] areas = (x2 - x1 + 1) * (y2 - y1 + 1) order = scores.argsort()[::-1] keep = [] while order.size > 0: i = order[0] keep.append(i) xx1 = np.maximum(x1[i], x1[order[1:]]) yy1 = np.maximum(y1[i], y1[order[1:]]) xx2 = np.minimum(x2[i], x2[order[1:]]) yy2 = np.minimum(y2[i], y2[order[1:]]) w = np.maximum(0.0, xx2 - xx1 + 1) h = np.maximum(0.0, yy2 - yy1 + 1) inter = w * h ovr = inter / (areas[i] + areas[order[1:]] - inter) inds = np.where(ovr <= IOU_THRESHOLD)[0] order = order[inds + 1] return keep

测试单张图像

def test_image(): PLUGIN_LIBRARY = "./libmyplugins.so" ctypes.CDLL(PLUGIN_LIBRARY) engine_file = "yolov5s.engine" image_file = "test.jpg" detector = YoLov5TRT(engine_file) img = cv2.imread(image_file) start = time.time() boxes, scores, class_ids = detector.infer(img) end = time.time() print(f"Inference time: {(end - start)*1000:.2f} ms") for i in range(len(boxes)): box = boxes[i] score = float(scores[i]) cls_id = int(class_ids[i]) plot_one_box(box, img, label=f"{categories[cls_id]} {score:.2f}") cv2.imwrite("output.jpg", img) print("Saved result to output.jpg")

运行命令：

python3 yolov5_trt.py

实时视频与摄像头推理

视频文件检测

def test_video(video_path): cap = cv2.VideoCapture(video_path) out = cv2.VideoWriter("output.avi", cv2.VideoWriter_fourcc(*'XVID'), 10, (int(cap.get(3)), int(cap.get(4)))) detector = YoLov5TRT("yolov5s.engine") while True: ret, frame = cap.read() if not ret: break boxes, scores, class_ids = detector.infer(frame) for i in range(len(boxes)): plot_one_box(boxes[i], frame, label=f"{categories[class_ids[i]]} {scores[i]:.2f}") out.write(frame) cv2.imshow("result", frame) if cv2.waitKey(1) == ord('q'): break cap.release() out.release() cv2.destroyAllWindows()

CSI摄像头实时推理（Jetson原生支持）

from jetcam.csi_camera import CSICamera def test_camera(): camera = CSICamera(width=640, height=480, capture_device=0) detector = YoLov5TRT("yolov5s.engine") while True: frame = camera.read() boxes, scores, class_ids = detector.infer(frame) for i in range(len(boxes)): plot_one_box(boxes[i], frame, label=f"{categories[class_ids[i]]} {scores[i]:.2f}") cv2.imshow("CSI Camera - YOLOv5 TRT", frame) if cv2.waitKey(1) == ord('q'): break camera.release() cv2.destroyAllWindows()

📌 安装jetcam：pip3 install jetcam

关于 YOLOv8 镜像的补充思考

如今已有预配置的“YOLO-V8 镜像”，内置PyTorch 2.x、Ultralytics框架、Jupyter Notebook等，适合快速原型开发。但对于生产部署，仍建议走TensorRT路线。

你可以：
1. 在高性能主机或云服务器上训练YOLOv8模型；
2. 导出为ONNX；
3. 移植至Jetson Nano并通过类似流程转为TensorRT引擎。

这样既能享受新模型带来的精度优势，又能获得极致推理效率。

常见问题与实战调试技巧

💡经验之谈：初次部署建议从yolov5n开始，即使精度略低，也能快速验证流程完整性。一旦跑通，再逐步升级模型规模。

这种将模型从训练域迁移到边缘推理域的过程，本质上是一场“瘦身革命”。它迫使开发者深入理解模型结构、TensorRT机制以及硬件限制之间的平衡。掌握这套方法论，不仅适用于YOLOv5，也为未来部署YOLOv8、YOLO-NAS乃至自定义网络打下坚实基础。

当你的Jetson Nano第一次流畅地识别出行人、车辆，并在屏幕上画出边界框时，那种“理论终于落地”的成就感，远胜于任何跑分数字。