1. YOLOv5模型ONNX导出与TensorRT部署基础

任务目标

本案例完成YOLOv5s模型从PyTorch权重到ONNX再到TensorRT引擎的完整转换链路，并基于转换后的TensorRT引擎执行一次基础推理，验证检测结果与原始PyTorch模型输出的一致性。

环境与依赖

  • Ubuntu 20.04+ / Windows 10+，Python 3.8+
  • PyTorch 1.12+，CUDA 11.6+，cuDNN 8.4+
  • TensorRT 8.4+（需与CUDA版本匹配）
  • YOLOv5官方仓库（v6.0+）：git clone https://github.com/ultralytics/yolov5
  • Python依赖：pip install onnx onnxruntime-gpu pycuda

实现步骤

1. 导出ONNX模型

在YOLOv5仓库根目录执行以下命令，将官方预训练权重yolov5s.pt导出为ONNX格式。导出时固定输入尺寸为640×640，并启用--simplify以优化计算图。

python export.py --weights yolov5s.pt --include onnx --img-size 640 --batch-size 1 --simplify

执行成功后生成yolov5s.onnx文件。可通过onnxruntime快速验证ONNX模型可正常加载并执行推理。

2. 构建TensorRT引擎

使用TensorRT的trtexec工具将ONNX模型转换为FP16精度的TensorRT引擎。FP16在保持精度的同时显著提升推理速度，适合自动驾驶场景的实时性要求。

trtexec --onnx=yolov5s.onnx --saveEngine=yolov5s_fp16.trt --fp16 --workspace=2048

转换完成后生成yolov5s_fp16.trt引擎文件。若需INT8量化，需额外提供校准数据集，本案例不展开。

3. TensorRT基础推理代码

以下代码实现TensorRT引擎的加载与推理。核心逻辑包括：创建运行时与引擎上下文、分配GPU显存、执行前向推理、将输出拷贝回CPU。输入为预处理后的640×640 RGB图像（归一化至[0,1]），输出为YOLOv5标准的3个检测头原始张量（形状分别为[1, 3, 80, 80, 85]、[1, 3, 40, 40, 85]、[1, 3, 20, 20, 85]）。

import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

class TRTInference:
    def __init__(self, engine_path):
        # 加载引擎
        with open(engine_path, 'rb') as f:
            engine_data = f.read()
        runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
        self.engine = runtime.deserialize_cuda_engine(engine_data)
        self.context = self.engine.create_execution_context()
        
        # 分配输入输出缓冲区
        self.inputs, self.outputs, self.bindings = [], [], []
        for binding in self.engine:
            shape = self.engine.get_binding_shape(binding)
            size = trt.volume(shape)
            dtype = trt.nptype(self.engine.get_binding_dtype(binding))
            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(binding):
                self.inputs.append({'name': binding, 'host': host_mem, 'device': device_mem, 'shape': shape})
            else:
                self.outputs.append({'name': binding, 'host': host_mem, 'device': device_mem, 'shape': shape})
    
    def infer(self, input_image):
        # 输入预处理：HWC->CHW，归一化
        img = input_image.astype(np.float32) / 255.0
        img = np.transpose(img, (2, 0, 1))[None]
        np.copyto(self.inputs[0]['host'], img.ravel())
        
        # 执行推理
        cuda.memcpy_htod(self.inputs[0]['device'], self.inputs[0]['host'])
        self.context.execute_v2(bindings=self.bindings)
        for out in self.outputs:
            cuda.memcpy_dtoh(out['host'], out['device'])
        
        # 返回三个检测头的输出
        return [out['host'].reshape(out['shape']) for out in self.outputs]

# 使用示例
engine_path = 'yolov5s_fp16.trt'
trt_infer = TRTInference(engine_path)

# 加载一张测试图像（需自行准备，此处仅演示接口）
# image = cv2.imread('test.jpg')
# image = cv2.resize(image, (640, 640))
# outputs = trt_infer.infer(image)
# print(f"检测头输出形状: {[o.shape for o in outputs]}")

4. 验证检测结果

验证分为两个层面：

  1. 输出一致性：使用同一张测试图像分别通过PyTorch模型（yolov5s.pt）和TensorRT引擎推理，对比输出张量的余弦相似度。余弦相似度大于0.99即认为转换精度损失可接受。
  2. 端到端检测：对TensorRT输出执行YOLOv5标准的NMS后处理（需复用YOLOv5仓库utils/general.py中的non_max_suppression函数），在图像上绘制检测框，确认目标被正确框出。

通过条件：TensorRT推理输出的检测框坐标与类别置信度与PyTorch模型结果基本一致（IoU > 0.9），且单帧推理延迟满足实时性要求（FP16下通常< 10ms，具体取决于GPU型号）。

常见问题

  • 导出ONNX时维度错误：确保--batch-size 1，动态batch需额外配置--dynamic参数。
  • TensorRT转换失败：检查ONNX是否包含不支持的操作（如部分上采样层），可尝试升级TensorRT版本或使用--fp16而非INT8。
  • 推理结果全零：检查输入图像是否已归一化至[0,1]且通道顺序为RGB（非BGR）。
← 上一章 📖 返回目录 下一章 →