TensorRT、OpenVINO、ONNXRuntime INT8量化：三份脚本，一张表，直接封神

# TensorRT INT8 量化示例脚本 (基于PyTorch + ModelOpt)
import torch
import modelopt.torch.quantization as mtq
import torch_tensorrt
# 1. 准备模型并设置为评估模式
model = MyModel().eval().cuda()
# 2. 定义校准数据加载循环 (使用你的真实数据)
def calibration_loop(model):
    for batch in calibration_dataloader:  # calibration_dataloader 需提前定义
        model(batch.cuda())
# 3. 应用INT8量化配置并进行校准
quant_cfg = mtq.INT8_DEFAULT_CFG  # 使用默认INT8配置
mtq.quantize(model, quant_cfg, forward_loop=calibration_loop)
# 4. 使用Torch-TensorRT编译生成INT8引擎
inputs = [torch.randn(1, 3, 224, 224).cuda()]  # 示例输入
trt_model = torch_tensorrt.compile(
    model,
    ir="dynamo",
    arg_inputs=inputs,
    min_block_size=1,
)
print("TensorRT INT8引擎生成完成！")

# OpenVINO INT8 量化示例脚本 (使用NNCF)

import nncf
from openvino.runtime import Core, serialize
from openvino.tools import mo
# 1. 加载FP32的OpenVINO IR模型
core = Core()
model = core.read_model("model_fp32.xml")
# 2. 准备校准数据集 (通常是一个小批量的DataLoader)
# 假设 validation_loader 是已有的数据加载器，这里仅作示意
def transform_fn(data_item):
    images, _ = data_item
    return images.numpy()
calibration_dataset = nncf.Dataset(validation_loader, transform_fn)
# 3. 应用INT8量化 (使用默认预设)
quantized_model = nncf.quantize(model, calibration_dataset, preset=nncf.QuantizationPreset.MIXED)
# 4. 保存量化后的INT8 IR模型
serialize(quantized_model, "model_int8.xml", "model_int8.bin")
print("OpenVINO INT8模型生成完成！")

# ONNX Runtime INT8 量化示例脚本
import onnx
from onnxruntime.quantization import quantize_static, QuantType, QuantFormat
from onnxruntime.quantization import CalibrationDataReader
# 1. 准备校准数据读取器 (需自定义以适配你的数据)
class MyDataReader(CalibrationDataReader):
    def __init__(self):
        super().__init__()
        # 初始化你的校准数据列表 (numpy数组)
        self.data = [{"input": np.random.randn(1, 3, 224, 224).astype(np.float32)} for _ in range(100)]
        self.iter = iter(self.data)
    def get_next(self):
        return next(self.iter, None)
# 2. 执行静态量化
input_model_path = "model_fp32.onnx"
output_model_path = "model_int8.onnx"
quantized_model = quantize_static(
    input_model_path,
    output_model_path,
    calibration_data_reader=MyDataReader(),
    quant_format=QuantFormat.QDQ,  # 使用QDQ格式
    activation_type=QuantType.QUInt8,
    weight_type=QuantType.QUInt8,
    per_channel=False,
    op_types_to_quantize=["Conv", "MatMul"],  # 指定需要量化的算子类型
)
print("ONNX Runtime INT8模型生成完成！")