2. YOLOv8n模型TensorRT FP16转换

任务目标

将YOLOv8n的ONNX模型转换为TensorRT FP16推理引擎，并验证转换后引擎的推理速度与正确性。本案例聚焦于TensorRT引擎构建与基础性能测试，不涉及端到端目标检测后处理逻辑。

环境依赖

• TensorRT 8.6+（含trtexec工具）
• CUDA 11.8+、cuDNN 8.6+
• Python 3.8+，安装pycuda、numpy、onnx
• 已导出的YOLOv8n ONNX模型（输入尺寸1x3x640x640，动态或固定batch均可）

实现步骤

整个流程分为两步：先用trtexec命令行完成ONNX到FP16引擎的转换，再通过Python API加载引擎并测试推理延迟。转换命令如下：

trtexec --onnx=yolov8n.onnx \
        --saveEngine=yolov8n_fp16.engine \
        --fp16 \
        --minShapes=images:1x3x640x640 \
        --optShapes=images:1x3x640x640 \
        --maxShapes=images:1x3x640x640

若ONNX模型已固定batch为1，可省略--minShapes/--optShapes/--maxShapes参数。转换成功后生成yolov8n_fp16.engine文件。

关键完整代码

以下Python代码加载生成的FP16引擎，对随机输入执行100次推理并统计平均延迟：

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

class TRTEngine:
    def __init__(self, engine_path):
        logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, 'rb') as f:
            runtime = trt.Runtime(logger)
            self.engine = runtime.deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()
        self._allocate_buffers()

    def _allocate_buffers(self):
        self.inputs, self.outputs, self.bindings = [], [], []
        for i in range(self.engine.num_bindings):
            name = self.engine.get_binding_name(i)
            shape = self.engine.get_binding_shape(i)
            size = trt.volume(shape)
            dtype = trt.nptype(self.engine.get_binding_dtype(i))
            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(name):
                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, input_data):
        inp = self.inputs[0]
        np.copyto(inp['host'], input_data.ravel())
        cuda.memcpy_htod(inp['device'], inp['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]

    def __del__(self):
        for binding in self.inputs + self.outputs:
            binding['device'].free()

engine = TRTEngine('yolov8n_fp16.engine')
input_data = np.random.randn(1, 3, 640, 640).astype(np.float32)

# 预热
for _ in range(10):
    engine.infer(input_data)

# 计时测试
times = []
for _ in range(100):
    start = time.perf_counter()
    engine.infer(input_data)
    times.append((time.perf_counter() - start) * 1000)

avg_ms = np.mean(times)
print(f"FP16平均推理延迟: {avg_ms:.2f} ms")
print(f"输出张量数量: {len(engine.outputs)}")
for i, out in enumerate(engine.outputs):
    print(f"输出{i} shape: {out['shape']}")

运行与验证

1. 确保yolov8n_fp16.engine与Python脚本在同一目录。
2. 执行脚本：python test_trt.py
3. 通过条件：程序正常结束，输出平均延迟为有限正数（通常小于30ms），且输出张量shape与ONNX模型输出一致（YOLOv8n输出为1x84x8400）。

若需对比FP32性能，可用trtexec --onnx=yolov8n.onnx --saveEngine=yolov8n_fp32.engine生成FP32引擎，替换引擎路径后重新运行脚本，对比两者平均延迟即可验证FP16加速效果。