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

# TensorRT INT8 量化示例腳本 (基于PyTorch + ModelOpt)import torchimport modelopt.torch.quantization as mtqimport torch_tensorrt# 1. 準(zhǔn)備模型并設(shè)置為評估模式model = MyModel().eval().cuda()# 2. 定義校準(zhǔn)數(shù)據(jù)加載循環(huán) (使用你的真實(shí)數(shù)據(jù))def calibration_loop(model):    for batch in calibration_dataloader:  # calibration_dataloader 需提前定義        model(batch.cuda())# 3. 應(yīng)用INT8量化配置并進(jìn)行校準(zhǔn)quant_cfg = mtq.INT8_DEFAULT_CFG  # 使用默認(rèn)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 nncffrom openvino.runtime import Core, serializefrom openvino.tools import mo# 1. 加載FP32的OpenVINO IR模型core = Core()model = core.read_model("model_fp32.xml")# 2. 準(zhǔn)備校準(zhǔn)數(shù)據(jù)集 (通常是一個小批量的DataLoader)# 假設(shè) validation_loader 是已有的數(shù)據(jù)加載器，這里僅作示意def transform_fn(data_item):    images, _ = data_item    return images.numpy()calibration_dataset = nncf.Dataset(validation_loader, transform_fn)# 3. 應(yīng)用INT8量化 (使用默認(rèn)預(yù)設(shè))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 onnxfrom onnxruntime.quantization import quantize_static, QuantType, QuantFormatfrom onnxruntime.quantization import CalibrationDataReader# 1. 準(zhǔn)備校準(zhǔn)數(shù)據(jù)讀取器 (需自定義以適配你的數(shù)據(jù))class MyDataReader(CalibrationDataReader):    def __init__(self):        super().__init__()        # 初始化你的校準(zhǔn)數(shù)據(jù)列表 (numpy數(shù)組)        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. 執(zhí)行靜態(tài)量化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模型生成完成！")