2. YOLOv5目标检测的TensorRT加速

任务目标

将训练好的YOLOv5s PyTorch模型转换为TensorRT FP16引擎，在嵌入式GPU（如Jetson Orin/Xavier）上完成目标检测推理，并对比TensorRT与PyTorch的端到端延迟与吞吐量。

环境与依赖

  • 硬件：Jetson Orin Nano/NX（或其他嵌入式GPU，算力≥7.5）
  • 软件：JetPack 5.x（含CUDA 11.4+、cuDNN 8.x）、TensorRT 8.5+、PyTorch 1.12+（仅用于导出与对比）
  • Python库：onnx、onnx-simplifier、pycuda、numpy、opencv-python
  • YOLOv5官方仓库（v6.0+）：git clone https://github.com/ultralytics/yolov5

实现步骤

整体流程：PyTorch权重 → ONNX（含NMS后处理）→ TensorRT FP16引擎 → 推理验证。本案例使用YOLOv5自带的NMS集成导出方式，避免在TensorRT中手动实现NMS解码。

1. 导出带NMS的ONNX模型

在YOLOv5仓库目录下执行以下命令，导出包含高效NMS插件的ONNX模型（输出维度为[1, 300, 6]，每行格式为[x1, y1, x2, y2, conf, class]）：

python export.py --weights yolov5s.pt --include onnx --opset 12 --simplify --dynamic False --conf-thres 0.25 --iou-thres 0.45 --max-det 300

关键参数说明：--dynamic False固定输入尺寸为640×640，避免TensorRT动态shape带来的额外显存开销；--conf-thres与--iou-thres在导出时固化到NMS节点中，推理时无需再设置。

2. 构建TensorRT FP16引擎

以下代码读取ONNX并构建FP16引擎，保存为.engine文件供后续直接加载推理：

import tensorrt as trt
import os

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def build_engine(onnx_path, engine_path, fp16=True):
    """将ONNX转换为TensorRT引擎，启用FP16"""
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, TRT_LOGGER)

    with open(onnx_path, 'rb') as f:
        if not parser.parse(f.read()):
            for i in range(parser.num_errors):
                print(parser.get_error(i))
            raise RuntimeError("ONNX解析失败")

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

    if fp16 and builder.platform_has_fast_fp16:
        config.set_flag(trt.BuilderFlag.FP16)
        print("已启用FP16")

    # 固定batch size=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)
    if engine is None:
        raise RuntimeError("引擎构建失败")

    with open(engine_path, 'wb') as f:
        f.write(engine)
    print(f"引擎已保存至: {engine_path}")

if __name__ == '__main__':
    build_engine('yolov5s.onnx', 'yolov5s_fp16.engine')

3. TensorRT推理与性能对比

加载引擎执行推理，并与PyTorch原始模型对比延迟。以下代码实现完整的推理管线与计时逻辑：

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import cv2
import time
import torch

# ---------- TensorRT 推理类 ----------
class TRTInference:
    def __init__(self, engine_path):
        with open(engine_path, 'rb') as f:
            runtime = trt.Runtime(TRT_LOGGER)
            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_io_tensors):
            name = self.engine.get_tensor_name(i)
            shape = self.engine.get_tensor_shape(name)
            dtype = trt.nptype(self.engine.get_tensor_dtype(name))
            size = trt.volume(shape)
            host_mem = cuda.pagelocked_empty(size, dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)
            self.bindings.append(int(device_mem))
            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                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})

    def infer(self, img_tensor):
        """img_tensor: 已归一化并转为NCHW的float32数组，shape=(1,3,640,640)"""
        np.copyto(self.inputs[0]['host'], img_tensor.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()
        # 输出shape: (1, 300, 6)
        return self.outputs[0]['host'].reshape(1, 300, 6)

# ---------- 预处理 ----------
def preprocess(img_path):
    img = cv2.imread(img_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (640, 640))
    img = img.astype(np.float32) / 255.0
    img = np.transpose(img, (2, 0, 1))  # HWC -> CHW
    return np.expand_dims(img, axis=0).copy()

# ---------- 性能对比 ----------
def benchmark():
    # 加载TensorRT引擎
    trt_engine = TRTInference('yolov5s_fp16.engine')
    img_tensor = preprocess('test.jpg')

    # Warm-up
    for _ in range(10):
        trt_engine.infer(img_tensor)

    # TensorRT 延迟测试（100次）
    trt_times = []
    for _ in range(100):
        t0 = time.perf_counter()
        trt_engine.infer(img_tensor)
        trt_times.append((time.perf_counter() - t0) * 1000)
    trt_avg = np.mean(trt_times[10:])  # 跳过前10次预热

    # PyTorch 对比（需先加载yolov5s.pt）
    model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s.pt', force_reload=True)
    model.eval().cuda()
    img_cv = cv2.imread('test.jpg')
    img_rgb = cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB)

    torch_times = []
    with torch.no_grad():
        for _ in range(10):  # 预热
            model(img_rgb, size=640)
        for _ in range(100):
            t0 = time.perf_counter()
            model(img_rgb, size=640)
            torch_times.append((time.perf_counter() - t0) * 1000)
    torch_avg = np.mean(torch_times[10:])

    print(f"TensorRT FP16 平均延迟: {trt_avg:.2f} ms")
    print(f"PyTorch FP32 平均延迟: {torch_avg:.2f} ms")
    print(f"加速比: {torch_avg / trt_avg:.2f}x")

if __name__ == '__main__':
    benchmark()

4. 验证检测结果

运行推理并绘制检测框，确认TensorRT输出与PyTorch结果一致（IoU > 0.5且类别相同）：

def verify_detection():
    trt_engine = TRTInference('yolov5s_fp16.engine')
    img_tensor = preprocess('test.jpg')
    dets = trt_engine.infer(img_tensor)[0]

    # 过滤置信度>0.25的检测框
    valid = dets[dets[:, 4] > 0.25]
    print(f"检测到 {len(valid)} 个目标")
    for det in valid:
        x1, y1, x2, y2, conf, cls = det
        print(f"类别{int(cls)} 置信度{conf:.2f} 框[{x1:.0f},{y1:.0f},{x2:.0f},{y2:.0f}]")

    # 与PyTorch结果对比（此处仅打印，实际可计算IoU）
    model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s.pt', force_reload=True)
    results = model(cv2.cvtColor(cv2.imread('test.jpg'), cv2.COLOR_BGR2RGB), size=640)
    print("\nPyTorch检测结果:")
    results.print()

if __name__ == '__main__':
    verify_detection()

运行与验证

  1. 准备一张包含常见物体的测试图片test.jpg（建议使用COCO类别中的物体，如人、车、狗）。
  2. 按顺序执行：先运行build_engine.py生成引擎，再运行benchmark()获取性能数据。
  3. 通过条件：
  • 引擎构建无报错，FP16成功启用；
  • TensorRT与PyTorch对同一张图检测到的目标类别一致，且置信度差异小于0.1；
  • TensorRT FP16延迟显著低于PyTorch FP32（在Jetson Orin上通常加速2～4倍）；
  • 单次推理延迟稳定，无显存溢出或崩溃。

常见问题

  • 引擎构建失败：确认ONNX模型已用onnx-simplifier简化，且opset≥12；若显存不足，将workspace从1GB降至512MB。
  • 输出全零：检查输入预处理是否与训练时一致（RGB顺序、归一化到0～1）；确认导出ONNX时的conf-thres未设置过高。
  • FP16精度下降：若检测框偏移明显，可回退到FP32引擎对比，或使用--iou-thres 0.5重新导出以降低NMS敏感度。
← 上一章 📖 返回目录 下一章 →