4. TensorRT FP16精度优化

任务目标

在已具备YOLO模型ONNX导出的基础上，使用TensorRT Python API分别构建FP32与FP16引擎，对同一测试图像执行推理，对比两种精度下的推理耗时与检测精度（mAP或单图IoU）。

环境依赖

• TensorRT 8.x+（含Python绑定）
• CUDA 11.x+、cuDNN 8.x+
• OpenCV-Python、NumPy
• 已导出的YOLOv5s ONNX模型（输入640×640，输出1×25200×85）

关键实现

以下代码完成FP32/FP16双引擎构建、推理耗时统计与精度对比。核心逻辑：build_engine()通过builder.create_network()解析ONNX，FP16模式需显式设置set_flag(trt.BuilderFlag.FP16)；推理阶段使用同一CUDA流，以消除流切换开销差异。

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

class EngineBuilder:
    def __init__(self, onnx_path, engine_path, fp16=False):
        self.onnx_path = onnx_path
        self.engine_path = engine_path
        self.fp16 = fp16
        self.logger = trt.Logger(trt.Logger.WARNING)

    def build(self):
        if os.path.exists(self.engine_path):
            return
        builder = trt.Builder(self.logger)
        network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
        parser = trt.OnnxParser(network, self.logger)
        with open(self.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)
        if self.fp16:
            config.set_flag(trt.BuilderFlag.FP16)
        # 固定batch=1
        profile = builder.create_optimization_profile()
        inp = network.get_input(0)
        profile.set_shape(inp.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(self.engine_path, 'wb') as f:
            f.write(engine)
        print(f"引擎已保存: {self.engine_path}")

def load_engine(engine_path):
    logger = trt.Logger(trt.Logger.WARNING)
    with open(engine_path, 'rb') as f:
        runtime = trt.Runtime(logger)
        return runtime.deserialize_cuda_engine(f.read())

def infer(engine, img_path, warmup=10, repeats=50):
    ctx = engine.create_execution_context()
    inp_name = engine.get_tensor_name(0)
    out_name = engine.get_tensor_name(1)
    inp_shape = (1, 3, 640, 640)
    out_shape = (1, 25200, 85)
    ctx.set_input_shape(inp_name, inp_shape)

    img = cv2.imread(img_path)
    img = cv2.resize(img, (640, 640))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
    img = img.transpose(2, 0, 1)[None]
    inp_data = np.ascontiguousarray(img)

    out_data = np.empty(out_shape, dtype=np.float32)
    d_in = cuda.mem_alloc(inp_data.nbytes)
    d_out = cuda.mem_alloc(out_data.nbytes)
    stream = cuda.Stream()

    # 拷贝输入
    cuda.memcpy_htod_async(d_in, inp_data, stream)
    ctx.set_tensor_address(inp_name, int(d_in))
    ctx.set_tensor_address(out_name, int(d_out))

    # 预热
    for _ in range(warmup):
        ctx.execute_async_v3(stream.handle)
        cuda.stream.synchronize(stream)

    # 计时推理
    times = []
    for _ in range(repeats):
        t0 = time.perf_counter()
        ctx.execute_async_v3(stream.handle)
        cuda.stream.synchronize(stream)
        times.append(time.perf_counter() - t0)

    cuda.memcpy_dtoh_async(out_data, d_out, stream)
    cuda.stream.synchronize(stream)
    d_in.free(); d_out.free()
    return out_data, np.mean(times[5:]) * 1000  # 去掉前5次抖动，返回ms

# 后处理：NMS + 计算IoU（与FP32/FP16对比用）
def postprocess(raw_output, conf_thres=0.25, iou_thres=0.45):
    boxes = raw_output[0]
    valid = boxes[:, 4] > conf_thres
    boxes = boxes[valid]
    if len(boxes) == 0:
        return []
    class_ids = np.argmax(boxes[:, 5:], axis=1)
    scores = boxes[np.arange(len(boxes)), 5 + class_ids]
    xywh = boxes[:, :4]
    xyxy = np.zeros_like(xywh)
    xyxy[:, 0] = xywh[:, 0] - xywh[:, 2] / 2
    xyxy[:, 1] = xywh[:, 1] - xywh[:, 3] / 2
    xyxy[:, 2] = xywh[:, 0] + xywh[:, 2] / 2
    xyxy[:, 3] = xywh[:, 1] + xywh[:, 3] / 2
    # 简单NMS
    indices = cv2.dnn.NMSBoxes(xyxy.tolist(), scores.tolist(), conf_thres, iou_thres)
    if len(indices) == 0:
        return []
    indices = indices.flatten()
    return [(xyxy[i], scores[i], class_ids[i]) for i in indices]

def compute_iou(box1, box2):
    x1 = max(box1[0], box2[0]); y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2]); y2 = min(box1[3], box2[3])
    inter = max(0, x2 - x1) * max(0, y2 - y1)
    area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
    area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
    return inter / (area1 + area2 - inter + 1e-6)

if __name__ == "__main__":
    import pycuda.driver as cuda
    import pycuda.autoinit

    # 1. 构建双引擎
    EngineBuilder("yolov5s.onnx", "yolov5s_fp32.engine", fp16=False).build()
    EngineBuilder("yolov5s.onnx", "yolov5s_fp16.engine", fp16=True).build()

    # 2. 加载并推理
    eng32 = load_engine("yolov5s_fp32.engine")
    eng16 = load_engine("yolov5s_fp16.engine")
    out32, t32 = infer(eng32, "test.jpg")
    out16, t16 = infer(eng16, "test.jpg")

    # 3. 精度对比（以FP32结果为基准）
    dets32 = postprocess(out32)
    dets16 = postprocess(out16)
    if len(dets32) == 0 or len(dets16) == 0:
        print("警告：某精度下无检测结果，请降低conf_thres")
    else:
        ious = []
        for d16 in dets16:
            best_iou = max(compute_iou(d16[0], d32[0]) for d32 in dets32)
            ious.append(best_iou)
        mean_iou = np.mean(ious)
        print(f"FP32 平均耗时: {t32:.2f} ms")
        print(f"FP16 平均耗时: {t16:.2f} ms")
        print(f"速度提升: {(t32/t16 - 1)*100:.1f}%")
        print(f"FP16相对FP32的检测框平均IoU: {mean_iou:.4f}")

    # 4. 释放资源
    del eng32, eng16

运行与验证

1. 安装依赖：pip install tensorrt pycuda opencv-python numpy
2. 将ONNX模型与测试图像置于脚本同目录，执行python trt_fp16_compare.py
3. 观察输出：应显示FP16较FP32有显著速度提升（通常1.5～2倍），且平均IoU > 0.9视为精度无损

通过标准：程序正常结束并打印四行指标；FP16耗时低于FP32；平均IoU大于0.85（若低于此值需检查是否因FP16动态范围导致输出溢出，可尝试在ONNX导出时对输出层使用FP32精度保留）。