将静态训练后量化(PTQ)应用于预训练的Transformer模型,需要使用常用库并遵循特定步骤。此过程的目标是将模型的权重和可能的激活计算转换为使用低精度整数(如INT8),通过校准数据集来确定最佳量化参数。请注意,此过程在模型完全训练之后进行。我们将主要使用Hugging Face生态系统,具体包括 transformers 库来加载模型和分词器,datasets 库来获取校准数据,以及 optimum 库,它提供模型优化的工具,包括量化,通常会借助ONNX Runtime等后端。设置您的环境首先,确保您已安装所需的库。您需要 transformers、datasets、optimum,以及 optimum 支持的量化后端,例如 onnxruntime。您可能还需要 accelerate 来更顺畅地处理模型。pip install transformers datasets optimum[onnxruntime] accelerate torch注意:请确保也安装了PyTorch,因为 transformers 和 optimum 经常需要它。加载预训练模型和分词器我们首先加载一个标准预训练模型。为了方便此示例,我们使用一个较小、广为人知的模型,如 distilbert-base-uncased,以保持过程易于管理。from transformers import AutoModelForSequenceClassification, AutoTokenizer model_id = "distilbert-base-uncased" # 加载适合特定任务的模型,例如序列分类 model = AutoModelForSequenceClassification.from_pretrained(model_id) tokenizer = AutoTokenizer.from_pretrained(model_id) print("原始模型已加载:") print(model)这会加载模型的标准FP32(32位浮点)版本。我们的目标是将其转换为静态INT8表示。准备校准数据集静态PTQ依赖于一个小型且有代表性的数据集,即校准数据集。该数据用于在推理期间观察模型内部激活值的分布。这些观察到的范围(最小值和最大值)对于计算将FP32值映射到INT8所需的缩放因子和零点非常重要。让我们加载数据集的一个小部分,例如“sst2”(斯坦福情感树库)数据集,它常用于分类任务。我们只需要几百个样本进行校准。from datasets import load_dataset # 加载适合模型任务的数据集(例如,情感分析) calibration_dataset_name = "sst2" num_calibration_samples = 200 # 少量样本通常足够 # 加载并选择子集 full_calibration_dataset = load_dataset(calibration_dataset_name, split="train") calibration_indices = list(range(num_calibration_samples)) calibration_subset = full_calibration_dataset.select(calibration_indices) # 预处理函数 def preprocess_function(examples): # 如果文本列名不同,请调整为实际的列名 'sentence' return tokenizer(examples["sentence"], padding="max_length", max_length=128, truncation=True) # 应用预处理 processed_calibration_dataset = calibration_subset.map(preprocess_function, batched=True) print(f"已准备好包含 {len(processed_calibration_dataset)} 个样本的校准数据集。") # 显示一个已处理样本的键 print("已处理样本的键:", processed_calibration_dataset[0].keys())关键在于,此数据集应反映模型在实际推理时将遇到的数据类型。我们使用模型的tokenizer对其进行预处理,确保输入与模型的预期相符。使用Optimum配置静态量化现在,我们使用 optimum 来定义我们的量化配置。我们指定需要静态INT8量化。Optimum 在底层使用外部工具包(例如本例中的ONNX Runtime的量化工具)。from optimum.onnxruntime import ORTQuantizer, AutoQuantizationConfig from optimum.onnxruntime.configuration import Arm64QuantizationConfig, AutoCalibrationConfig # 1. 定义量化策略:静态INT8 qconfig = AutoQuantizationConfig.int8_static( config=model.config, # 提供模型配置 dataset=processed_calibration_dataset # 之前准备好的校准数据集 ) # 对于ONNX Runtime,我们通常先导出到ONNX onnx_model_path = "distilbert_base_uncased_onnx" quantized_model_path = "distilbert_base_uncased_quantized_onnx" # 2. 使用模型ID(或路径)创建量化器 quantizer = ORTQuantizer.from_pretrained(model_id, feature=model.config.task_specific_params.get("feature", "sequence-classification")) print("量化配置已创建:") print(qconfig)在这里,AutoQuantizationConfig.int8_static 方便地设置了静态INT8量化的配置。我们提供了模型的配置和我们已处理的校准数据集。ORTQuantizer 使用模型标识符进行初始化。应用静态PTQ过程模型加载完毕,校准数据准备就绪,配置也已定义,我们现在可以执行量化。quantizer.quantize 方法处理将模型导出为中间格式(如ONNX)、运行校准、计算量化参数,并生成最终的量化模型。# 3. 应用量化 quantizer.quantize( save_dir=quantized_model_path, quantization_config=qconfig, # 可选:指定中间ONNX模型的路径 # onnx_model_path=onnx_model_path ) print(f"静态PTQ完成。量化模型已保存到:{quantized_model_path}")此过程可能需要几分钟,具体取决于模型大小和校准样本数量。它包括:将原始FP32模型导出为ONNX格式。使用ONNX模型在校准数据集上运行推理,以收集激活统计信息(最小值/最大值范围)。使用这些统计信息来计算权重和激活的INT8缩放因子和零点。将模型权重转换为INT8,并修改计算图,在必要时加入量化和反量化节点。保存最终的量化ONNX模型。验证量化模型(大小比较)量化的一个主要好处是模型大小的减小。让我们比较原始PyTorch模型的状态字典大小与量化后的ONNX模型文件大小。import os # 获取目录大小的函数 def get_dir_size(path='.'): total = 0 with os.scandir(path) as it: for entry in it: if entry.is_file(): total += entry.stat().st_size elif entry.is_dir(): total += get_dir_size(entry.path) return total / (1024 * 1024) # 大小,单位MB # 原始PyTorch模型的近似大小(可能略有不同) original_model_size_mb = sum(p.numel() * p.element_size() for p in model.parameters()) / (1024 * 1024) # 量化ONNX模型目录的大小 quantized_model_size_mb = get_dir_size(quantized_model_path) print(f"原始FP32模型近似大小:{original_model_size_mb:.2f} MB") print(f"量化INT8模型大小(ONNX):{quantized_model_size_mb:.2f} MB") # 计算缩减量 reduction = (1 - (quantized_model_size_mb / original_model_size_mb)) * 100 print(f"大小缩减:{reduction:.2f}%") # 可选:可视化大小比较{"data": [{"type": "bar", "x": ["原始FP32", "量化INT8"], "y": [256.36, 67.5], "marker": {"color": ["#4263eb", "#12b886"]}}], "layout": {"title": "模型大小比较(近似值)", "yaxis": {"title": "大小 (MB)"}}}原始FP32模型与静态量化INT8模型(ONNX格式)之间的近似大小比较。大小为指示性数值,可能因具体模型和保存格式而异。您应该会观察到模型大小显著减小,从FP32到INT8通常接近4倍,因为INT8每个参数只使用四分之一的比特。加载并运行量化模型(可选)尽管详细评估将在第六章介绍,但您可以使用 optimum 和ONNX Runtime加载量化模型,进行快速推理检查。# 需要安装onnxruntime from optimum.onnxruntime import ORTModelForSequenceClassification from transformers import pipeline # 加载量化模型 quantized_ort_model = ORTModelForSequenceClassification.from_pretrained(quantized_model_path) # 创建一个方便推理的管道 classifier = pipeline("sentiment-analysis", model=quantized_ort_model, tokenizer=tokenizer) # 测试推理 text = "This movie was quite good, I enjoyed it." result = classifier(text) print("量化模型的推理结果:", result) text_neg = "The plot was predictable and the acting was mediocre." result_neg = classifier(text_neg) print("量化模型的推理结果:", result_neg)这确认量化模型可以加载并生成输出。请记住,静态PTQ有时会导致精度下降,这需要针对您的具体用例进行仔细评估。总结在本实践部分中,我们使用 optimum 库对 transformers 模型应用了静态训练后量化。我们加载了一个预训练模型,准备了校准数据集,配置了静态INT8转换的量化过程,使用ONNX Runtime作为后端执行了量化,并观察了由此产生的模型大小减小。这个实用的流程展示了PTQ如何有效应用,无需访问原始训练流程或昂贵的再训练,从而使模型更小并可能加快推理速度。接下来的章节将介绍旨在减轻潜在精度损失的更高级PTQ技术,并介绍量化感知训练(QAT)。