4. TensorRT引擎转换与推理

任务目标

将训练好的YOLOv5s ONNX模型（FP32）转换为TensorRT FP16引擎，并实现推理。通过对比TensorRT与ONNX Runtime的推理结果和耗时，验证转换正确性及加速效果。

环境依赖

• TensorRT 8.6+（含trtexec命令行工具）
• CUDA 11.8+、cuDNN 8.6+
• Python 3.8+，安装：pycuda、onnxruntime-gpu、numpy、opencv-python
• YOLOv5s ONNX模型文件（输入尺寸640×640，含NMS后处理部分）

实现方案

使用TensorRT Python API完成ONNX→Engine的转换，采用FP16精度以提升速度。推理阶段手动管理GPU显存，将预处理后的图像输入引擎，获取输出后执行置信度过滤与NMS。精度验证采用同一张测试图片，对比TensorRT与ONNX Runtime的检测框IoU及类别一致性。

完整代码

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

# ---------- 1. ONNX → TensorRT引擎转换 ----------
def build_engine(onnx_path, engine_path, fp16=True):
    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
    if fp16 and builder.platform_has_fast_fp16:
        config.set_flag(trt.BuilderFlag.FP16)
    
    # 固定batch size为1
    profile = builder.create_optimization_profile()
    input_tensor = network.get_input(0)
    input_shape = input_tensor.shape  # [1,3,640,640]
    profile.set_shape(input_tensor.name, input_shape, input_shape, input_shape)
    config.add_optimization_profile(profile)
    
    engine = builder.build_serialized_network(network, config)
    with open(engine_path, 'wb') as f:
        f.write(engine)
    return engine_path

# ---------- 2. TensorRT推理类 ----------
class TRTInference:
    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.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, host_mem, device_mem, shape))
            else:
                self.outputs.append((name, host_mem, device_mem, shape))
    
    def infer(self, input_data):
        # input_data: [1,3,640,640] float32, 已归一化
        inp_name, inp_host, inp_dev, inp_shape = self.inputs[0]
        np.copyto(inp_host, input_data.ravel())
        cuda.memcpy_htod(inp_dev, inp_host)
        
        self.context.set_tensor_address(inp_name, inp_dev)
        for out_name, out_host, out_dev, _ in self.outputs:
            self.context.set_tensor_address(out_name, out_dev)
        
        self.context.execute_async_v3(stream.handle)
        for _, out_host, out_dev, _ in self.outputs:
            cuda.memcpy_dtoh(out_host, out_dev)
        cuda.Stream.synchronize(stream)
        
        results = []
        for _, out_host, _, out_shape in self.outputs:
            results.append(out_host.reshape(out_shape).copy())
        return results

# ---------- 3. 后处理（YOLOv5输出格式） ----------
def postprocess(outputs, conf_thres=0.5, iou_thres=0.45):
    # outputs[0]: [1, 25200, 85] (cx,cy,w,h,obj_conf,cls_prob...)
    pred = outputs[0][0]
    scores = pred[:, 4] * pred[:, 5:].max(axis=1)
    mask = scores > conf_thres
    pred = pred[mask]
    scores = scores[mask]
    if len(pred) == 0:
        return []
    
    boxes = pred[:, :4].copy()
    # xywh → xyxy
    boxes[:, 0] -= boxes[:, 2] / 2
    boxes[:, 1] -= boxes[:, 3] / 2
    boxes[:, 2] += boxes[:, 0]
    boxes[:, 3] += boxes[:, 1]
    
    cls_ids = pred[:, 5:].argmax(axis=1)
    # 简单NMS（按类别）
    keep = []
    for c in np.unique(cls_ids):
        idx = np.where(cls_ids == c)[0]
        if len(idx) == 0:
            continue
        order = scores[idx].argsort()[::-1]
        idx = idx[order]
        while len(idx) > 0:
            keep.append(idx[0])
            if len(idx) == 1:
                break
            iou = compute_iou(boxes[idx[0]], boxes[idx[1:]])
            idx = idx[1:][iou < iou_thres]
    return [(boxes[i], scores[i], cls_ids[i]) for i in keep]

def compute_iou(box, boxes):
    x1 = np.maximum(box[0], boxes[:, 0])
    y1 = np.maximum(box[1], boxes[:, 1])
    x2 = np.minimum(box[2], boxes[:, 2])
    y2 = np.minimum(box[3], boxes[:, 3])
    inter = np.maximum(0, x2-x1) * np.maximum(0, y2-y1)
    area1 = (box[2]-box[0]) * (box[3]-box[1])
    area2 = (boxes[:,2]-boxes[:,0]) * (boxes[:,3]-boxes[:,1])
    return inter / (area1 + area2 - inter + 1e-6)

# ---------- 4. 主流程 ----------
if __name__ == '__main__':
    # 初始化CUDA流
    stream = cuda.Stream()
    
    # 转换引擎（首次运行）
    onnx_path = 'yolov5s.onnx'
    engine_path = 'yolov5s_fp16.engine'
    build_engine(onnx_path, engine_path, fp16=True)
    
    # 加载TensorRT引擎
    trt_infer = TRTInference(engine_path)
    
    # 加载ONNX Runtime（用于精度对比）
    ort_session = ort.InferenceSession(onnx_path, providers=['CUDAExecutionProvider'])
    
    # 读取并预处理图像
    img = cv2.imread('test.jpg')
    img_resized = cv2.resize(img, (640, 640))
    img_rgb = cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB)
    img_norm = img_rgb.astype(np.float32) / 255.0
    input_tensor = np.transpose(img_norm, (2, 0, 1))[None]  # [1,3,640,640]
    
    # TensorRT推理
    trt_start = time.time()
    trt_outputs = trt_infer.infer(input_tensor)
    trt_time = time.time() - trt_start
    trt_dets = postprocess(trt_outputs)
    
    # ONNX Runtime推理
    ort_start = time.time()
    ort_outputs = ort_session.run(None, {ort_session.get_inputs()[0].name: input_tensor})
    ort_time = time.time() - ort_start
    ort_dets = postprocess(ort_outputs)
    
    # 验证：比较检测框数量与类别
    print(f"TensorRT检测框数: {len(trt_dets)}, 耗时: {trt_time*1000:.2f}ms")
    print(f"ONNX Runtime检测框数: {len(ort_dets)}, 耗时: {ort_time*1000:.2f}ms")
    print(f"加速比: {ort_time/trt_time:.2f}x")
    
    # 逐框比较（按置信度排序后匹配）
    trt_dets.sort(key=lambda x: x[1], reverse=True)
    ort_dets.sort(key=lambda x: x[1], reverse=True)
    min_len = min(len(trt_dets), len(ort_dets))
    match_count = 0
    for i in range(min_len):
        trt_box, trt_score, trt_cls = trt_dets[i]
        ort_box, ort_score, ort_cls = ort_dets[i]
        if trt_cls == ort_cls:
            iou = compute_iou(trt_box, np.array([ort_box]))[0]
            if iou > 0.5:
                match_count += 1
    print(f"匹配框数(类别一致且IoU>0.5): {match_count}/{min_len}")

运行与验证

1. 将YOLOv5s ONNX模型和测试图片放在脚本同目录，修改代码中onnx_path和img路径。
2. 执行脚本：python trt_infer.py。首次运行会生成yolov5s_fp16.engine文件（约需1-2分钟）。
3. 观察输出：TensorRT与ONNX Runtime的检测框数量应一致或接近（允许置信度阈值边缘差异）。
4. 判定通过条件：匹配框数 ≥ min_len的80%，且TensorRT推理耗时显著低于ONNX Runtime（通常快2-5倍）。

常见问题

• 显存不足：降低WORKSPACE大小至512MB，或使用batch size=1的profile。
• FP16精度损失：若检测框匹配率低于80%，可改用FP32引擎（设置fp16=False）重新转换。
• 输出维度不匹配：确认ONNX模型已包含NMS输出层（输出形状为[1,25200,85]），若为原始模型需先导出带后处理的版本。