2. YOLOv5模型移植到TensorRT进行实时缺陷检测

任务目标

将训练好的YOLOv5s缺陷检测模型（PyTorch权重）转换为TensorRT FP16引擎，在NVIDIA GPU上实现实时推理，并验证检测精度与速度满足工业缺陷检测要求。

环境依赖

Ubuntu 20.04+ / Windows 10+，NVIDIA GPU（Compute Capability ≥ 7.5）
CUDA 11.4+，cuDNN 8.2+，TensorRT 8.6+
Python 3.8+，PyTorch 1.12+，OpenCV 4.5+
YOLOv5官方仓库（v6.0+，含export.py）

实现步骤

核心链路为：PyTorch权重 → ONNX → TensorRT FP16引擎 → 推理验证。使用YOLOv5官方export.py完成ONNX导出，再通过TensorRT Python API构建引擎并执行推理。

1. 导出ONNX模型

# 在YOLOv5仓库目录下执行
python export.py --weights best.pt --include onnx --opset 12 --simplify --dynamic
# 生成 best.onnx，确认输入尺寸为 [1, 3, 640, 640]

2. 构建TensorRT引擎并推理（关键完整代码）

import tensorrt as trt
import numpy as np
import cv2
import time

class YOLOv5TensorRT:
    def __init__(self, onnx_path, engine_path, fp16=True):
        self.logger = trt.Logger(trt.Logger.WARNING)
        self.engine = self.build_engine(onnx_path, engine_path, fp16)
        self.context = self.engine.create_execution_context()
        self.inputs, self.outputs, self.bindings = [], [], []
        self.stream = None
        self._allocate_buffers()

    def build_engine(self, onnx_path, engine_path, fp16):
        # 若引擎已存在则直接加载
        import os
        if os.path.exists(engine_path):
            with open(engine_path, 'rb') as f:
                return trt.Runtime(self.logger).deserialize_cuda_engine(f.read())
        # 否则构建
        builder = trt.Builder(self.logger)
        network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
        parser = trt.OnnxParser(network, self.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
        if fp16 and builder.platform_has_fast_fp16:
            config.set_flag(trt.BuilderFlag.FP16)
        # 设置动态batch
        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(engine_path, 'wb') as f:
            f.write(engine)
        return trt.Runtime(self.logger).deserialize_cuda_engine(engine)

    def _allocate_buffers(self):
        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 = np.empty(size, dtype=dtype)
            device_mem = None  # 使用tensorrt的cuda内存管理
            self.bindings.append(device_mem)
            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                self.inputs.append({'name': name, 'host': host_mem, 'shape': shape})
            else:
                self.outputs.append({'name': name, 'host': host_mem, 'shape': shape})

    def infer(self, image):
        # 预处理：resize + letterbox + normalize
        img, ratio, (dw, dh) = self.letterbox(image, (640, 640))
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR→RGB, HWC→CHW
        img = np.ascontiguousarray(img / 255.0, dtype=np.float32)
        # 设置输入
        self.context.set_input_shape(self.inputs[0]['name'], (1, 3, 640, 640))
        self.inputs[0]['host'] = img.ravel()
        # 执行推理（同步）
        self.context.execute_v2([self.inputs[0]['host']] + [out['host'] for out in self.outputs])
        # 后处理（简化：取第一个输出，shape为[1, 25200, 85]）
        pred = self.outputs[0]['host'].reshape(1, 25200, 85)
        boxes = self.postprocess(pred[0], ratio, (dw, dh))
        return boxes

    def letterbox(self, img, new_shape=(640, 640)):
        shape = img.shape[:2]
        r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
        new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r)))
        dw, dh = (new_shape[1] - new_unpad[0]) / 2, (new_shape[0] - new_unpad[1]) / 2
        img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
        top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
        left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
        img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
        return img, r, (dw, dh)

    def postprocess(self, pred, ratio, (dw, dh)):
        # 简化后处理：仅做置信度过滤和坐标还原（完整NMS可接入torchvision或自定义）
        conf_thres = 0.25
        xywh = pred[:, :4].copy()
        xyxy = np.column_stack([xywh[:, 0] - xywh[:, 2] / 2, xywh[:, 1] - xywh[:, 3] / 2,
                                xywh[:, 0] + xywh[:, 2] / 2, xywh[:, 1] + xywh[:, 3] / 2])
        scores = pred[:, 4:].max(axis=1)
        mask = scores > conf_thres
        if not mask.any():
            return []
        boxes = xyxy[mask]
        boxes[:, [0, 2]] = (boxes[:, [0, 2]] - dw) / ratio
        boxes[:, [1, 3]] = (boxes[:, [1, 3]] - dh) / ratio
        return boxes.astype(int).tolist()

# 使用示例
if __name__ == '__main__':
    detector = YOLOv5TensorRT('best.onnx', 'best_fp16.engine', fp16=True)
    img = cv2.imread('defect_sample.jpg')
    # 预热
    for _ in range(10):
        detector.infer(img)
    # 测速
    times = []
    for _ in range(100):
        t0 = time.time()
        boxes = detector.infer(img)
        times.append(time.time() - t0)
    print(f"平均推理耗时: {np.mean(times)*1000:.2f} ms")
    print(f"检测到缺陷框: {len(boxes)} 个")

运行与验证

将best.pt放在YOLOv5仓库根目录，执行导出命令生成best.onnx。
将上述代码保存为trt_infer.py，与best.onnx和测试图片defect_sample.jpg置于同一目录。
运行python trt_infer.py，观察输出。

通过条件：

程序无报错，成功加载引擎并输出平均推理耗时（FP16下通常< 5ms @ 640×640，具体取决于GPU型号）。
检测框数量与PyTorch原始模型在相同置信度阈值下的输出一致（允许坐标微小浮点差异）。
若需验证精度，可对同一测试集分别运行PyTorch模型和TensorRT引擎，计算mAP@0.5差异应< 0.5%。

常见问题

ONNX导出失败：确认PyTorch版本≥1.12，且--simplify需要安装onnx-simplifier。
引擎构建报错：检查GPU算力是否≥7.5，CUDA/cuDNN版本与TensorRT匹配。
输出shape不匹配：若模型输入尺寸非640，需同步修改代码中的letterbox目标尺寸和reshape参数。

📚 课程目录