3. YOLOX模型TensorRT FP16量化部署

任务目标

将训练好的YOLOX模型（以YOLOX-S为例）转换为TensorRT FP16推理引擎，在保持检测精度基本不降的前提下，获得显著的推理加速。本案例聚焦于完整的转换流程、精度对比方法和加速比测量，不涉及模型重新训练。

环境依赖

• TensorRT 8.x+（含trtexec工具）
• CUDA 11.x + cuDNN 8.x
• PyTorch 1.10+（用于导出ONNX与精度验证）
• YOLOX官方代码库（提供模型定义与权重）
• OpenCV、pycuda（用于Python推理脚本）

实现步骤与关键代码

第一步：导出ONNX模型。使用YOLOX官方tools/export_onnx.py导出不带NMS的ONNX文件，输入尺寸固定为640×640。导出后可用onnx-simplifier简化计算图。

# 导出命令（在YOLOX仓库根目录执行）
python tools/export_onnx.py --name yolox-s --ckpt yolox_s.pth --input 640 --output-name yolox_s.onnx --no-onnxsim
# 简化ONNX（可选但推荐）
python -m onnxsim yolox_s.onnx yolox_s_sim.onnx

第二步：转换为TensorRT FP16引擎。使用trtexec命令行工具完成转换，关键参数为FP16精度与动态batch支持。

# 转换为FP16引擎（静态batch=1）
trtexec --onnx=yolox_s_sim.onnx --saveEngine=yolox_s_fp16.engine --fp16 --workspace=4096

# 如需动态batch，使用以下参数
trtexec --onnx=yolox_s_sim.onnx --saveEngine=yolox_s_fp16_dyn.engine --fp16 --minShapes=images:1x3x640x640 --optShapes=images:1x3x640x640 --maxShapes=images:8x3x640x640

第三步：编写Python推理脚本。以下代码实现TensorRT引擎加载、预处理、推理与后处理（解码+NMS），用于后续精度与速度验证。

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

class YOLOX_TRT:
    def __init__(self, engine_path):
        self.logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, 'rb') as f:
            runtime = trt.Runtime(self.logger)
            self.engine = runtime.deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()
        self.stream = cuda.Stream()
        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(i):
                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, image):
        # 预处理：resize到640x640，归一化到[0,1]，CHW格式
        img = cv2.resize(image, (640, 640))
        img = img.astype(np.float32) / 255.0
        img = img.transpose(2, 0, 1)[None]
        np.copyto(self.inputs[0]['host'], img.ravel())
        cuda.memcpy_htod_async(self.inputs[0]['device'], self.inputs[0]['host'], self.stream)
        self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
        for out in self.outputs:
            cuda.memcpy_dtoh_async(out['host'], out['device'], self.stream)
        self.stream.synchronize()
        # 输出形状: [1, 8400, 85] (cx,cy,w,h,obj_conf,cls_scores)
        output = self.outputs[0]['host'].reshape(1, 8400, 85)
        return output

    def postprocess(self, output, conf_thres=0.5, iou_thres=0.45):
        # 解码：将cx,cy,w,h转为x1,y1,x2,y2
        boxes = output[0, :, :4].copy()
        boxes[:, 0] -= boxes[:, 2] / 2
        boxes[:, 1] -= boxes[:, 3] / 2
        boxes[:, 2] += boxes[:, 0]
        boxes[:, 3] += boxes[:, 1]
        scores = output[0, :, 4:]
        class_ids = np.argmax(scores, axis=1)
        confs = scores[np.arange(len(scores)), class_ids]
        mask = confs > conf_thres
        boxes, confs, class_ids = boxes[mask], confs[mask], class_ids[mask]
        if len(boxes) == 0:
            return [], [], []
        # 简单NMS（按类别）
        keep = []
        for cls in np.unique(class_ids):
            idx = np.where(class_ids == cls)[0]
            boxes_cls, confs_cls = boxes[idx], confs[idx]
            x1, y1, x2, y2 = boxes_cls[:, 0], boxes_cls[:, 1], boxes_cls[:, 2], boxes_cls[:, 3]
            areas = (x2 - x1) * (y2 - y1)
            order = confs_cls.argsort()[::-1]
            while order.size > 0:
                i = order[0]
                keep.append(idx[i])
                xx1 = np.maximum(x1[i], x1[order[1:]])
                yy1 = np.maximum(y1[i], y1[order[1:]])
                xx2 = np.minimum(x2[i], x2[order[1:]])
                yy2 = np.minimum(y2[i], y2[order[1:]])
                inter = np.maximum(0, xx2 - xx1) * np.maximum(0, yy2 - yy1)
                iou = inter / (areas[i] + areas[order[1:]] - inter)
                order = order[1:][iou <= iou_thres]
        return boxes[keep], confs[keep], class_ids[keep]

    def __del__(self):
        for inp in self.inputs:
            cuda.mem_free(inp['device'])
        for out in self.outputs:
            cuda.mem_free(out['device'])

第四步：精度与速度评估脚本。使用COCO验证集子集（或自定义测试图片）对比FP32 PyTorch模型与FP16 TensorRT模型的mAP，并测量推理延迟。

import torch
import time
from yolox.exp import get_exp
from yolox.utils import postprocess as yolox_postprocess

# 加载PyTorch FP32模型
exp = get_exp(None, 'yolox-s')
model = exp.get_model()
ckpt = torch.load('yolox_s.pth', map_location='cpu')
model.load_state_dict(ckpt['model'])
model.eval().cuda()

# 加载TensorRT引擎
trt_model = YOLOX_TRT('yolox_s_fp16.engine')

# 精度对比：在100张验证图片上计算mAP（简化版，仅统计IoU>0.5的匹配数）
def evaluate_map(model_type):
    tp, fp, num_gt = 0, 0, 0
    for img_path in val_images[:100]:
        img = cv2.imread(img_path)
        h, w = img.shape[:2]
        if model_type == 'pytorch':
            with torch.no_grad():
                outputs = model(torch.from_numpy(img.transpose(2,0,1)[None].astype(np.float32)/255).cuda())
            outputs = yolox_postprocess(outputs, 0.5, 0.45)[0]
            if outputs is None: continue
            boxes = outputs[:, :4].cpu().numpy()
            boxes[:, [0,2]] *= w/640; boxes[:, [1,3]] *= h/640
            scores = outputs[:, 4].cpu().numpy()
        else:
            out = trt_model.infer(img)
            boxes, scores, _ = trt_model.postprocess(out)
            boxes[:, [0,2]] *= w/640; boxes[:, [1,3]] *= h/640
        # 与GT匹配（此处省略详细匹配逻辑，实际需加载GT标注）
        # 统计tp/fp/num_gt...
    return tp / (tp + fp + 1e-6)  # 简化返回precision

# 延迟测量（预热后测100次平均）
def measure_latency(model_type, n=100):
    img = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
    if model_type == 'pytorch':
        for _ in range(10):  # 预热
            with torch.no_grad():
                model(torch.from_numpy(img.transpose(2,0,1)[None].astype(np.float32)/255).cuda())
        torch.cuda.synchronize()
        start = time.time()
        for _ in range(n):
            with torch.no_grad():
                model(torch.from_numpy(img.transpose(2,0,1)[None].astype(np.float32)/255).cuda())
        torch.cuda.synchronize()
    else:
        for _ in range(10):
            trt_model.infer(img)
        start = time.time()
        for _ in range(n):
            trt_model.infer(img)
    return (time.time() - start) / n * 1000  # ms

# 执行评估
pytorch_map = evaluate_map('pytorch')
trt_map = evaluate_map('tensorrt')
pytorch_lat = measure_latency('pytorch')
trt_lat = measure_latency('tensorrt')
print(f'PyTorch FP32: mAP={pytorch_map:.4f}, 延迟={pytorch_lat:.2f}ms')
print(f'TensorRT FP16: mAP={trt_map:.4f}, 延迟={trt_lat:.2f}ms')
print(f'加速比: {pytorch_lat/trt_lat:.2f}x, 精度损失: {pytorch_map-trt_map:.4f}')

验证方法与通过标准

1. 引擎生成验证：trtexec命令执行后无报错，生成yolox_s_fp16.engine文件且大小合理（约25-30MB）。
2. 功能正确性：运行推理脚本，对单张测试图片（含已知目标）能正确输出检测框，且框位置与PyTorch FP32结果基本一致（IoU>0.7）。
3. 精度评估：在100张验证图片上，FP16引擎的mAP与FP32 PyTorch模型相比，下降不超过1%（例如FP32 mAP=0.452，FP16 mAP≥0.442）。
4. 加速效果：FP16引擎的端到端推理延迟（含预处理与后处理）应显著低于PyTorch FP32，加速比≥2x（在相同GPU上测量）。