YOLOX模型Int8量化转TensorRT引擎失败问题求助

问题根源 #

• TensorRT对ONNX Runtime生成的QDQ格式模型兼容性有限，报错中的DequantizeLinear节点作用于bias张量，而TensorRT仅允许该层处理激活类型的输入，不支持直接量化bias。
• 使用随机dummy数据做校准，数据分布与真实场景差异大，不仅会导致量化精度下降，还可能引发模型结构解析异常。

可行方案 #

方案1：使用TensorRT原生静态量化（推荐） #

跳过ONNX Runtime的量化步骤，直接用TensorRT的原生工具完成Int8校准与引擎构建，兼容性最佳。

关键代码实现

首先实现自定义校准器（基于真实数据集）：

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

class YOLOXCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, calibration_images_dir, input_shape, batch_size=1, cache_file="calibration.cache"):
        super().__init__()
        self.cache_file = cache_file
        self.batch_size = batch_size
        self.input_shape = input_shape  # (3, 640, 640)
        self.image_paths = [os.path.join(calibration_images_dir, f) for f in os.listdir(calibration_images_dir) if f.endswith(('jpg', 'png'))]
        self.current_index = 0

        # 分配CUDA内存
        self.device_input = cuda.mem_alloc(trt.volume(input_shape) * batch_size * np.dtype(np.float32).itemsize)

    def get_batch_size(self):
        return self.batch_size

    def get_batch(self, names):
        if self.current_index + self.batch_size > len(self.image_paths):
            return None
        
        batch_data = []
        for i in range(self.batch_size):
            img_path = self.image_paths[self.current_index + i]
            img = cv2.imread(img_path)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = cv2.resize(img, (self.input_shape[1], self.input_shape[2]))
            img = img.transpose((2, 0, 1)).astype(np.float32)
            img /= 255.0  # 匹配YOLOX输入预处理逻辑
            batch_data.append(img)
        
        batch_data = np.stack(batch_data)
        cuda.memcpy_htod(self.device_input, batch_data.ravel())
        self.current_index += self.batch_size
        return [self.device_input]

    def read_calibration_cache(self):
        if os.path.exists(self.cache_file):
            with open(self.cache_file, "rb") as f:
                return f.read()
        return None

    def write_calibration_cache(self, cache):
        with open(self.cache_file, "wb") as f:
            f.write(cache)

然后构建Int8引擎：

TRT_LOGGER = trt.Logger(trt.Logger.INFO)
EXPLICIT_BATCH = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)

def build_int8_engine(onnx_path, engine_path, input_shape, max_workspace_size=30, batch_size=1, calibration_dir="calibration_images"):
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as network, builder.create_builder_config() as config, trt.OnnxParser(network, TRT_LOGGER) as parser, trt.Runtime(TRT_LOGGER) as runtime:
        config.max_workspace_size = 1 << max_workspace_size  # 30代表1GB显存
        builder.max_batch_size = batch_size

        # 开启Int8量化模式
        config.set_flag(trt.BuilderFlag.INT8)
        calibrator = YOLOXCalibrator(calibration_dir, input_shape[1:], batch_size=batch_size)
        config.int8_calibrator = calibrator

        # 解析ONNX模型
        with open(onnx_path, 'rb') as model:
            if not parser.parse(model.read()):
                for error in range(parser.num_errors):
                    print(parser.get_error(error))
                return None
        
        network.get_input(0).shape = input_shape
        plan = builder.build_serialized_network(network, config)
        engine = runtime.deserialize_cuda_engine(plan)
        
        with open(engine_path, "wb") as f:
            f.write(plan)
        return engine

# 使用示例
input_shape = (1, 3, 640, 640)
build_int8_engine("yolox_l.onnx", "yolox_l_int8.engine", input_shape, calibration_dir="path/to/your/calibration_images")

方案2：调整ONNX Runtime量化参数，规避不兼容节点 #

修改量化配置，排除bias节点的量化，同时使用真实校准数据：

修改后的量化代码

import numpy as np
import cv2
import os
import onnx
from onnxruntime.quantization import quantize_static, CalibrationMethod, CalibrationDataReader, QuantType, QuantFormat

# 基于真实数据的校准器
class RealDataReader(CalibrationDataReader):
    def __init__(self, image_dir, num_samples, input_shape):
        self.image_dir = image_dir
        self.num_samples = num_samples
        self.input_shape = input_shape
        self.current_sample = 0
        self.image_paths = [os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.endswith(('jpg', 'png'))][:num_samples]

    def get_next(self):
        if self.current_sample < self.num_samples:
            img_path = self.image_paths[self.current_sample]
            img = cv2.imread(img_path)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = cv2.resize(img, (self.input_shape[2], self.input_shape[3]))
            img = img.transpose((2, 0, 1)).astype(np.float32)
            img /= 255.0
            self.current_sample += 1
            return {'images': img[np.newaxis, ...]}
        else:
            return None

# 量化配置
onnx_model_input_path = "yolox_l.onnx"
onnx_model_output_path = "yolox_l_qdq_int8.onnx"
num_calibration_samples = 100
input_shape = (1, 3, 640, 640)

# 自动排除所有含bias的节点
model = onnx.load(onnx_model_input_path)
excluded_nodes = [node.name for node in model.graph.node if "bias" in node.name.lower()]

calibration_data_reader = RealDataReader("path/to/calibration_images", num_calibration_samples, input_shape)

quantized_model = quantize_static(
    model_input=onnx_model_input_path,
    model_output=onnx_model_output_path,
    calibration_data_reader=calibration_data_reader,
    activation_type=QuantType.QInt8,
    weight_type=QuantType.QInt8,
    quant_format=QuantFormat.QDQ,
    per_channel=True,  # 权重按通道量化，提升精度与兼容性
    calibrate_method=CalibrationMethod.Entropy,
    excluded_nodes=excluded_nodes  # 跳过bias节点的量化
)

转换时使用原TensorRT构建代码，但无需开启Int8模式（模型已为量化格式）。

方案3：使用YOLOX官方TensorRT量化脚本 #

YOLOX官方提供了直接导出TensorRT Int8引擎的工具，无需手动处理ONNX量化：

1. 克隆YOLOX仓库并安装依赖
2. 运行官方脚本完成导出与量化：

# 先导出float32 ONNX模型
python tools/export_onnx.py --output-name yolox_l.onnx -n yolox-l -c path/to/yolox_l.pth

# 直接生成Int8 TensorRT引擎
python tools/trt_quant.py --onnx yolox_l.onnx --engine yolox_l_int8.engine --calib-images path/to/calibration_images --batch-size 1

注意事项 #

• 校准数据集必须包含真实场景图像，建议500-1000张，确保量化后精度达标。
• 使用方案1时，需保证CUDA环境与TensorRT版本匹配（推荐TensorRT 8.x及以上）。
• 方案3需严格遵循YOLOX官方的环境配置要求。

内容的提问来源于stack exchange，提问作者codeine