ONNX-Python静态量化中Calibration_Data_Reader的创建与通用实现问询

核心问题 #

你把MinMaxCalibrater和CalibrationDataReader搞混了——这是两个完全不同的东西：

• Calibrater是ONNX Runtime内部用来统计数据分布的组件
• CalibrationDataReader是你必须自己实现的自定义类，作用是给校准器批量喂输入数据

如何实现CalibrationDataReader #

只需要给类实现一个get_next()方法，要求：

• 返回值是字典，key是模型输入节点的名称（可以用Netron这类工具查看模型结构获取）
• value是对应输入的numpy数组，形状、数据类型必须和模型输入完全匹配

通用实现模板 #

import numpy as np
from onnxruntime.quantization import quantize_static, CalibrationMethod, create_calibrator

class GenericCalibrationDataReader:
    def __init__(self, data_generator, input_names):
        self.data_generator = data_generator
        self.input_names = input_names

    def get_next(self):
        try:
            # 从数据生成器取一批数据
            batch_data = next(self.data_generator)
            # 把数据映射成模型要求的输入字典
            return {name: batch_data[i] for i, name in enumerate(self.input_names)}
        except StopIteration:
            # 返回None表示所有校准数据已读完
            return None

# 示例：替换成你的真实校准数据加载逻辑
def get_calibration_data(batch_size=8, num_batches=10):
    for _ in range(num_batches):
        # 假设模型有两个输入：input_1 (shape [8,3,224,224])、input_2 (shape [8,10])
        yield (
            np.random.randn(batch_size, 3, 224, 224).astype(np.float32),
            np.random.randint(0, 10, size=(batch_size,10)).astype(np.float32)
        )

修正你的量化代码 #

正确的流程是：创建数据读取器 → 用读取器给校准器喂数据 → 执行量化

def quantize_model(model_path, output_path, input_names):
    # 1. 初始化校准数据读取器
    data_gen = get_calibration_data()
    data_reader = GenericCalibrationDataReader(data_gen, input_names)

    # 2. 创建校准器并收集数据分布
    calibrator = create_calibrator(
        model_path,
        calibrate_method=CalibrationMethod.MinMax
    )
    calibrator.collect_data(data_reader)

    # 3. 执行静态量化
    quantize_static(
        model_input=model_path,
        model_output=output_path,
        calibration_data_reader=data_reader,
        calibrator=calibrator
    )

# 使用时替换成你的模型输入节点名称
quantize_model("original_model.onnx", "quantized_model.onnx", ["input_1", "input_2"])

通用读取器的可行性 #

完全可以做通用的CalibrationDataReader：

• 把数据加载逻辑抽成独立的生成器，读取器只负责将生成器输出映射到模型输入名
• 换模型时只需要传入新的input_names和对应的数据生成器即可适配

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