3. 基于TensorRT的YOLOv5目标检测部署

任务目标

在Jetson设备上，将PyTorch训练的YOLOv5s模型转换为TensorRT FP16引擎，并编写推理脚本实现USB摄像头实时目标检测，要求单帧预处理+推理+后处理总耗时低于30ms（Jetson Orin系列）。

环境依赖

Jetson设备（Orin Nano/NX/AGX），JetPack 5.1+（含CUDA 11.4、TensorRT 8.5）
Python 3.8+，PyTorch 1.12+（仅转换阶段需要）
安装依赖：pip install ultralytics opencv-python numpy pycuda
YOLOv5s.pt权重（ultralytics官方或自训练）

实现步骤

整体流程分为两步：离线转换（生成engine文件）与在线推理（加载engine执行检测）。转换阶段在PC或Jetson上均可完成，推理阶段必须在Jetson上运行。

1. 导出ONNX并转换为TensorRT引擎

使用ultralytics导出接口生成ONNX，再通过TensorRT的Python API构建FP16引擎。关键点：需显式指定opset=12以兼容TensorRT的GridSample层；动态batch设为固定1以降低显存占用。

# export_engine.py
from ultralytics import YOLO
import tensorrt as trt
import os

# 1. 导出ONNX
model = YOLO("yolov5su.pt")
model.export(format="onnx", opset=12, imgsz=640, dynamic=False, simplify=True)
onnx_path = "yolov5su.onnx"

# 2. 构建TensorRT引擎
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)

with open(onnx_path, "rb") as f:
    assert parser.parse(f.read()), "ONNX解析失败"

config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
config.set_flag(trt.BuilderFlag.FP16)

# 固定batch=1
profile = builder.create_optimization_profile()
input_tensor = network.get_input(0)
profile.set_shape(input_tensor.name, (1, 3, 640, 640), (1, 3, 640, 640), (1, 3, 640, 640))
config.add_optimization_profile(profile)

engine = builder.build_serialized_network(network, config)
with open("yolov5su.engine", "wb") as f:
    f.write(engine)
print("engine已生成")

2. Jetson实时推理脚本

推理脚本负责：加载engine → 创建CUDA上下文 → 摄像头采集 → 预处理（letterbox+归一化）→ 推理 → NMS后处理 → 绘制结果。输出为检测框坐标、类别与置信度。

# detect_live.py
import cv2
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

# 类别名称（COCO 80类）
CLASSES = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat',
           'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
           'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella',
           'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite',
           'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle',
           'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
           'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
           'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
           'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book',
           'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']

class YOLOv5TRT:
    def __init__(self, engine_path):
        logger = trt.Logger(trt.Logger.WARNING)
        runtime = trt.Runtime(logger)
        with open(engine_path, 'rb') as f:
            self.engine = runtime.deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()
        self.stream = cuda.Stream()
        
        # 分配输入输出缓冲区
        self.inputs, self.outputs, self.bindings = [], [], []
        for i in range(self.engine.num_bindings):
            name = self.engine.get_binding_name(i)
            shape = self.engine.get_binding_shape(i)
            size = trt.volume(shape)
            dtype = trt.nptype(self.engine.get_binding_dtype(i))
            host_mem = cuda.pagelocked_empty(size, dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)
            self.bindings.append(int(device_mem))
            if self.engine.binding_is_input(i):
                self.inputs.append({'name': name, 'host': host_mem, 'device': device_mem, 'shape': shape})
            else:
                self.outputs.append({'name': name, 'host': host_mem, 'device': device_mem, 'shape': shape})
        
        self.output_shape = self.outputs[0]['shape']  # (1, 25200, 85)

    def preprocess(self, img):
        # letterbox填充至640x640
        h, w = img.shape[:2]
        ratio = min(640 / w, 640 / h)
        new_w, new_h = int(w * ratio), int(h * ratio)
        resized = cv2.resize(img, (new_w, new_h))
        canvas = np.full((640, 640, 3), 114, dtype=np.uint8)
        canvas[:new_h, :new_w] = resized
        # BGR->RGB, HWC->CHW, 归一化
        blob = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB).transpose(2, 0, 1).astype(np.float32) / 255.0
        return np.ascontiguousarray(blob), ratio, (new_w, new_h)

    def infer(self, img):
        blob, ratio, (new_w, new_h) = self.preprocess(img)
        self.inputs[0]['host'][:] = blob.ravel()
        cuda.memcpy_htod_async(self.inputs[0]['device'], self.inputs[0]['host'], self.stream)
        self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
        cuda.memcpy_dtoh_async(self.outputs[0]['host'], self.outputs[0]['device'], self.stream)
        self.stream.synchronize()
        
        pred = self.outputs[0]['host'].reshape(self.output_shape)[0]  # (25200, 85)
        return self.postprocess(pred, ratio, (new_w, new_h), img.shape)

    def postprocess(self, pred, ratio, new_shape, orig_shape):
        # 过滤低置信度
        conf_mask = pred[:, 4] > 0.25
        pred = pred[conf_mask]
        if len(pred) == 0:
            return []
        
        boxes = pred[:, :4].copy()
        # 将cx,cy,w,h转换为x1,y1,x2,y2（letterbox坐标系）
        boxes[:, 0] = pred[:, 0] - pred[:, 2] / 2
        boxes[:, 1] = pred[:, 1] - pred[:, 3] / 2
        boxes[:, 2] = pred[:, 0] + pred[:, 2] / 2
        boxes[:, 3] = pred[:, 1] + pred[:, 3] / 2
        
        # 映射回原图坐标
        boxes[:, [0, 2]] /= ratio
        boxes[:, [1, 3]] /= ratio
        boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0, orig_shape[1])
        boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0, orig_shape[0])
        
        scores = pred[:, 4]
        cls_ids = np.argmax(pred[:, 5:], axis=1)
        cls_scores = pred[:, 5:].max(axis=1)
        final_scores = scores * cls_scores
        
        # 简单NMS
        indices = cv2.dnn.NMSBoxes(boxes.tolist(), final_scores.tolist(), 0.25, 0.45)
        results = []
        for i in indices.flatten():
            x1, y1, x2, y2 = boxes[i].astype(int)
            results.append((x1, y1, x2, y2, float(final_scores[i]), int(cls_ids[i])))
        return results

def main():
    model = YOLOv5TRT("yolov5su.engine")
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        dets = model.infer(frame)
        for x1, y1, x2, y2, score, cls_id in dets:
            label = f"{CLASSES[cls_id]} {score:.2f}"
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
        
        cv2.imshow("YOLOv5 TRT", frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()

运行与验证

生成引擎：在具备PyTorch的环境执行python export_engine.py，得到yolov5su.engine文件。
实时检测：将engine文件拷贝至Jetson，连接USB摄像头后执行python detect_live.py。
通过标准：画面窗口正常显示，检测框稳定跟随目标，帧率≥30FPS（可用cv2.getTickCount自行测量耗时），且检测类别与位置基本准确。

常见问题

engine构建失败：确认ONNX导出时已设置opset=12且simplify=True；若显存不足，将WORKSPACE降至512MB。
输出shape不匹配：YOLOv5s输出为(1, 25200, 85)，若使用其他模型变体需调整后处理中的维度参数。
首次运行较慢：TensorRT引擎加载需1-2秒，属正常现象；若持续低帧率，检查是否启用了FP16（Jetson Nano不支持）。

📚 课程目录