YOLOv5模型在Jetson设备上的TensorRT部署-开发者社区

YOLOv5模型在Jetson设备上的TensorRT部署

在边缘计算场景中，如何让一个训练好的目标检测模型真正“落地”运行，是每个AI工程师绕不开的课题。尤其是在NVIDIA Jetson Nano、Xavier NX、Orin这类资源受限但又需要实时响应的嵌入式平台上，性能与功耗的平衡变得异常关键。

YOLO（You Only Look Once）系列模型因其速度快、精度高，早已成为工业界首选的目标检测方案。虽然YOLOv5并非Ultralytics官方正式命名版本（实为社区广泛采用的实现），但它凭借简洁的代码结构、高效的训练流程和出色的推理表现，迅速成为实际项目中最常用的框架之一。

更令人振奋的是，如今我们已拥有成熟的工具链支持：从 PyTorch 模型导出，到 TensorRT 引擎编译，再到 Jetson 上 C++/Python 的端到端调用，整个部署路径已经非常清晰。本文将以YOLOv5 在 Jetson 系列设备上通过 TensorRT 加速推理为主线，手把手带你完成从模型导出到实时视频流处理的全流程，并结合当前流行的预装开发镜像进行适配说明。

环境准备：硬件选型与系统配置

要跑通这套部署流程，首先得有一块能打的“板子”。

软件环境要求

主流 Jetson 开发基于JetPack SDK，它集成了 L4T（Linux for Tegra）、CUDA、cuDNN、TensorRT 等核心组件。推荐使用以下版本组合：

组件	推荐版本
JetPack	5.1 或 5.0.2
CUDA	12.x / 11.4
cuDNN	9.x
TensorRT	8.6+
OpenCV	4.5+
Python	3.8 ~ 3.10

你可以通过以下命令查看当前系统的 L4T 版本信息：

sudo apt install lsb-release cat /etc/nv_tegra_release

确认 JetPack 版本后，再选择对应的依赖库版本，避免因版本错配导致构建失败。

模型导出：生成可用于 TensorRT 的中间文件

YOLOv5 提供了export.py工具，可以直接将.pt权重转换为 ONNX、TorchScript 或原生的TensorRT Engine (.engine)文件。这里介绍两种常用方式。

方法一：直接导出`.engine`文件（推荐）

如果你是在具备完整 CUDA/TensorRT 环境的主机上操作（比如带 GPU 的 x86 PC），可以一步到位生成优化后的引擎文件。

进入 YOLOv5 项目目录：

cd /root/ultralytics/yolov5

执行导出命令：

yolo export model=best.pt format=engine imgsz=640 device=0

参数说明：

model: 训练好的权重路径，如runs/train/exp/weights/best.pt
format=engine: 输出格式为 TensorRT 引擎
imgsz=640: 输入图像尺寸，需为 32 的倍数
device=0: 使用 GPU 编译引擎

成功后会生成best.engine文件，该文件已完成量化、层融合等优化，可直接拷贝至 Jetson 设备加载。

✅ 优势：无需手动构建插件，省去复杂编译步骤
❌ 缺点：必须在同一架构（aarch64）且有 GPU 的环境下运行，否则无法生成有效引擎

方法二：生成`.wts`文件（跨平台通用方案）

若你在 x86 主机训练模型，却要在 aarch64 架构的 Jetson 上部署，则应采用生成.wts文件 + 在 Jetson 上本地 build engine的方式。

步骤 1：生成`.wts`权重文件

下载 tensorrtx 仓库中的 YOLOv5 支持代码：

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

复制gen_wts.py到你的 YOLOv5 目录下：

cp tensorrtx/yolov5/gen_wts.py yolov5/ cd yolov5 python3 gen_wts.py -w best.pt

这会在当前目录生成best.wts文件，其中包含所有层名与对应权重，便于后续在 Jetson 上重建网络并插入参数。

然后将其传输到 Jetson 设备：

scp best.wts user@jetson-ip:/home/user/tensorrtx/yolov5/

Jetson 端部署：构建 TensorRT 引擎

现在我们切换到 Jetson 设备，开始真正的“本地化”构建过程。

安装必要依赖

确保已安装基础库和 Python 包：

sudo apt update sudo apt install python3-pip libopencv-dev python3-opencv libssl-dev pip3 install pycuda numpy

验证 TensorRT 是否可用：

import tensorrt as trt print(trt.__version__)

如果没有报错，说明环境正常。

编译 YOLOv5-TensorRT 示例程序

进入tensorrtx/yolov5目录：

cd ~/tensorrtx/yolov5

修改类别数量与输入尺寸

打开yololayer.h文件，调整以下宏定义：

#define CLASS_NUM 80 // 根据你的数据集类别数修改 #define INPUT_W 640 // 与训练时一致 #define INPUT_H 640

例如，如果你训练的是自定义四分类模型（人、车、狗、猫），则改为：

#define CLASS_NUM 4

⚠️ 务必保持与训练配置完全一致，否则输出解析将出错！

构建工程

mkdir build && cd build cmake .. make -j$(nproc)

如果提示找不到libmyplugins.so，请先执行一次-s参数来生成 engine 和插件库：

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

该命令含义如下：

# sudo ./yolov5 -s [.wts] [.engine] [model_type] # 示例： sudo ./yolov5 -s ../best.wts best.engine s

支持的模型类型包括：s,m,l,x,s6,m6,l6,x6，也可使用自定义宽度/深度因子：

sudo ./yolov5 -s ../custom.wts custom.engine c 0.33 0.50

其中c表示 custom，0.33是 depth_multiple，0.50是 width_multiple。

构建成功后你会得到两个关键产物：

best.engine：序列化的 TensorRT 推理引擎
libmyplugins.so：自定义插件库（含 YOLO 层实现）

这两个文件将在后续推理中被同时加载。

推理测试：C++ 与 Python 双模式调用

C++ 测试图片推理

将待测图片放入samples/文件夹：

cp ../data/images/bus.jpg samples/

运行推理：

sudo ./yolov5 -d best.engine ../samples

结果图片将保存为output*.jpg。若没有检测框输出，请检查：

CONF_THRESH和NMS_THRESH是否设置过高（默认 0.5 和 0.4）
类别数是否匹配
图片路径是否正确

Python 调用（更适合集成应用）

为了方便后续接入 ROS、Flask API 或摄像头流，推荐使用 Python 调用引擎。

编辑yolov5_trt.py文件，进行以下关键修改：

（1）禁用 Torch 相关模块

注释掉对torch和torchvision的引用，避免在 Jetson 上额外安装这些重型包：

# import torch # import torchvision

（2）统一输入尺寸

确保与训练和导出时一致：

INPUT_W = 640 INPUT_H = 640

（3）替换 NMS 实现为 NumPy 版本

删除对torchvision.ops.nms的调用，改用纯 NumPy 实现的非极大值抑制函数：

def nms(self, boxes, scores, iou_threshold=0.4): 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

并在post_process中调用它：

indices = self.nms(boxes, scores, IOU_THRESHOLD) result_boxes = boxes[indices] result_scores = scores[indices] result_classid = classid[indices]

（4）设置自定义类别标签

CLASSES = ['person', 'bicycle', 'car', 'motorcycle'] # 替换为你自己的类别

完整 Python 推理脚本示例

import cv2 import numpy as np import time import ctypes import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit # 加载自定义插件库 PLUGIN_LIBRARY = "libmyplugins.so" ctypes.CDLL(PLUGIN_LIBRARY) # 参数设置 ENGINE_PATH = "best.engine" IMAGE_PATH = "samples/bus.jpg" INPUT_W, INPUT_H = 640, 640 CONF_THRESH = 0.4 IOU_THRESHOLD = 0.5 CLASSES = ['person', 'bicycle', 'car', 'motorcycle'] class YOLOv5TRT: def __init__(self, engine_path): self.logger = trt.Logger(trt.Logger.INFO) with open(engine_path, "rb") as f: runtime = trt.Runtime(self.logger) self.engine = runtime.deserialize_cuda_engine(f.read()) self.context = self.engine.create_execution_context() self.stream = cuda.Stream() # 分配内存缓冲区 self.host_inputs = [] self.cuda_inputs = [] self.host_outputs = [] self.cuda_outputs = [] for binding in self.engine: size = trt.volume(self.engine.get_binding_shape(binding)) * self.engine.max_batch_size dtype = trt.nptype(self.engine.get_binding_dtype(binding)) host_mem = cuda.pagelocked_empty(size, dtype) cuda_mem = cuda.mem_alloc(host_mem.nbytes) if self.engine.binding_is_input(binding): self.host_inputs.append(host_mem) self.cuda_inputs.append(cuda_mem) else: self.host_outputs.append(host_mem) self.cuda_outputs.append(cuda_mem) def infer(self, image): # 预处理：BGR → RGB → resize → transpose → normalize input_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) input_image = cv2.resize(input_image, (INPUT_W, INPUT_H)) input_image = np.transpose(input_image, (2, 0, 1)).astype(np.float32) / 255.0 input_image = np.ascontiguousarray(input_image) np.copyto(self.host_inputs[0], input_image.ravel()) # Host to Device cuda.memcpy_htod_async(self.cuda_inputs[0], self.host_inputs[0], self.stream) # 执行异步推理 self.context.execute_async_v2( bindings=[int(d) for d in self.cuda_inputs + self.cuda_outputs], stream_handle=self.stream.handle ) # Device to Host cuda.memcpy_dtoh_async(self.host_outputs[0], self.cuda_outputs[0], self.stream) self.stream.synchronize() output = self.host_outputs[0] num_dets = int(output[0]) det_data = output[1:].reshape(-1, 6)[:num_dets] boxes = det_data[:, :4] scores = det_data[:, 4] class_ids = det_data[:, 5].astype(int) # 置信度过滤 conf_mask = scores > CONF_THRESH boxes = boxes[conf_mask] scores = scores[conf_mask] class_ids = class_ids[conf_mask] # 坐标还原到原始图像尺度 scale_x = image.shape[1] / INPUT_W scale_y = image.shape[0] / INPUT_H boxes[:, 0::2] *= scale_x boxes[:, 1::2] *= scale_y # NMS 后处理 indices = self.nms(boxes, scores, IOU_THRESHOLD) return boxes[indices], scores[indices], class_ids[indices] @staticmethod def nms(boxes, scores, threshold): 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 <= threshold)[0] order = order[inds + 1] return keep # 绘图函数 def draw_boxes(img, boxes, scores, class_ids, class_names): for i in range(len(boxes)): x1, y1, x2, y2 = map(int, boxes[i]) label = f"{class_names[class_ids[i]]}: {scores[i]:.2f}" color = (0, 255, 0) cv2.rectangle(img, (x1, y1), (x2, y2), color, 2) cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) return img # 主函数 if __name__ == "__main__": detector = YOLOv5TRT(ENGINE_PATH) image = cv2.imread(IMAGE_PATH) start = time.time() boxes, scores, class_ids = detector.infer(image) print(f"Inference time: {(time.time() - start)*1000:.2f} ms") result_img = draw_boxes(image, boxes, scores, class_ids, CLASSES) cv2.imwrite("result_detection.jpg", result_img) print("Detection result saved to result_detection.jpg")

运行后即可看到带有检测框的结果图。

视频与摄像头实时推理

只需稍作扩展，就能实现视频或摄像头实时检测。

视频文件推理

def detect_video(video_path): cap = cv2.VideoCapture(video_path) out = cv2.VideoWriter( "output.avi", cv2.VideoWriter_fourcc(*"MJPG"), 30, (int(cap.get(3)), int(cap.get(4))) ) while cap.isOpened(): ret, frame = cap.read() if not ret: break boxes, scores, class_ids = detector.infer(frame) result_frame = draw_boxes(frame, boxes, scores, class_ids, CLASSES) out.write(result_frame) cv2.imshow("YOLOv5-TensorRT", result_frame) if cv2.waitKey(1) == ord('q'): break cap.release() out.release() cv2.destroyAllWindows()

CSI 摄像头支持（Jetson 原生）

对于 Jetson 原生 CSI 摄像头（如 Raspberry Pi Camera V2），可通过jetcam库直接访问：

from jetcam.csi_camera import CSICamera camera = CSICamera(width=640, height=480, capture_fps=30) while True: frame = camera.read() boxes, scores, class_ids = detector.infer(frame) result_frame = draw_boxes(frame, boxes, scores, class_ids, CLASSES) cv2.imshow("CSI Camera", result_frame) if cv2.waitKey(1) == ord('q'): break cv2.destroyAllWindows()

这种方式延迟低、稳定性好，非常适合做边缘 AI 产品原型。

关于 YOLO-V8 镜像的特别说明

尽管本文聚焦于 YOLOv5，但现实中很多开发者使用的是预装好的YOLO-V8 开发镜像—— 这类镜像通常基于 Docker 或 SD 卡烧录形式分发，特点包括：

预装 PyTorch、Ultralytics、OpenCV、TensorRT
提供 Jupyter Notebook 开发界面
内置ultralytics项目目录
支持一键训练与导出

你可以在其中直接运行 YOLOv5 的训练与导出任务：

cd /root/ultralytics yolo task=detect mode=train model=yolov5s.yaml data=coco.yaml epochs=100 imgsz=640 yolo export model=runs/detect/train/weights/best.pt format=engine

再将生成的.engine或.wts文件拷出，用于 Jetson 部署。

这类镜像大大降低了环境配置门槛，尤其适合新手快速上手。不过要注意其内部 Python 环境可能与标准 JetPack 存在差异，在部署时仍需验证 TensorRT 兼容性。

整个流程走下来，你会发现虽然涉及不少底层细节——比如插件编译、内存管理、NMS 实现等——但一旦打通，就能在 Jetson 上实现10~30 FPS 的实时检测能力，足以支撑大多数智能安防、机器人导航、工业质检等边缘 AI 场景。

更重要的是，这种“训练-导出-部署”的标准化路径，也为后续升级到 YOLOv8、打包成 Flask API 或集成进 ROS 节点打下了坚实基础。下一步，不妨尝试把模型封装成 REST 接口，或是加入跟踪算法实现多目标追踪，你会发现，真正的 AI 落地之旅才刚刚开始。

YOLOv5模型在Jetson设备上的TensorRT部署