快速入门

🤗 Optimum 的核心是使用配置对象来定义不同加速器上的优化参数。然后，这些对象用于实例化专用的优化器、量化器和剪枝器。

在应用量化或优化之前，我们首先需要将我们的模型导出为 ONNX 格式。

>>> from optimum.onnxruntime import ORTModelForSequenceClassification
>>> from transformers import AutoTokenizer

>>> model_checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"
>>> save_directory = "tmp/onnx/"
>>> # Load a model from transformers and export it to ONNX
>>> ort_model = ORTModelForSequenceClassification.from_pretrained(model_checkpoint, export=True)
>>> tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
>>> # Save the onnx model and tokenizer
>>> ort_model.save_pretrained(save_directory)
>>> tokenizer.save_pretrained(save_directory)

现在让我们看看如何使用 ONNX Runtime 应用动态量化

>>> from optimum.onnxruntime.configuration import AutoQuantizationConfig
>>> from optimum.onnxruntime import ORTQuantizer
>>> # Define the quantization methodology
>>> qconfig = AutoQuantizationConfig.arm64(is_static=False, per_channel=False)
>>> quantizer = ORTQuantizer.from_pretrained(ort_model)
>>> # Apply dynamic quantization on the model
>>> quantizer.quantize(save_dir=save_directory, quantization_config=qconfig)

在本示例中，我们量化了来自 Hugging Face Hub 的模型，但它也可以是本地模型目录的路径。应用 quantize() 方法的结果是一个 model_quantized.onnx 文件，可用于运行推理。以下是如何加载 ONNX Runtime 模型并使用它生成预测的示例

>>> from optimum.onnxruntime import ORTModelForSequenceClassification
>>> from transformers import pipeline, AutoTokenizer
>>> model = ORTModelForSequenceClassification.from_pretrained(save_directory, file_name="model_quantized.onnx")
>>> tokenizer = AutoTokenizer.from_pretrained(save_directory)
>>> cls_pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)
>>> results = cls_pipeline("I love burritos!")

类似地，您可以通过在实例化 QuantizationConfig 对象时简单地将 is_static 设置为 True 来应用静态量化

>>> qconfig = AutoQuantizationConfig.arm64(is_static=True, per_channel=False)

静态量化依赖于将数据批次馈送到模型中，以在推理时间之前估计激活量化参数。为了支持这一点，🤗 Optimum 允许您提供校准数据集。校准数据集可以是来自 🤗 Datasets 库的简单 Dataset 对象，也可以是托管在 Hugging Face Hub 上的任何数据集。对于本示例，我们将选择模型最初训练时使用的 sst2 数据集

>>> from functools import partial
>>> from optimum.onnxruntime.configuration import AutoCalibrationConfig

# Define the processing function to apply to each example after loading the dataset
>>> def preprocess_fn(ex, tokenizer):
...     return tokenizer(ex["sentence"])

>>> # Create the calibration dataset
>>> calibration_dataset = quantizer.get_calibration_dataset(
...     "glue",
...     dataset_config_name="sst2",
...     preprocess_function=partial(preprocess_fn, tokenizer=tokenizer),
...     num_samples=50,
...     dataset_split="train",
... )
>>> # Create the calibration configuration containing the parameters related to calibration.
>>> calibration_config = AutoCalibrationConfig.minmax(calibration_dataset)
>>> # Perform the calibration step: computes the activations quantization ranges
>>> ranges = quantizer.fit(
...     dataset=calibration_dataset,
...     calibration_config=calibration_config,
...     operators_to_quantize=qconfig.operators_to_quantize,
... )
>>> # Apply static quantization on the model
>>> quantizer.quantize(
...     save_dir=save_directory,
...     calibration_tensors_range=ranges,
...     quantization_config=qconfig,
... )

作为最后一个示例，让我们看一下应用图优化技术，例如运算符融合和常量折叠。和以前一样，我们加载一个配置对象，但这次是通过设置优化级别而不是量化方法

>>> from optimum.onnxruntime.configuration import OptimizationConfig

>>> # Here the optimization level is selected to be 1, enabling basic optimizations such as redundant node eliminations and constant folding. Higher optimization level will result in a hardware dependent optimized graph.
>>> optimization_config = OptimizationConfig(optimization_level=1)

接下来，我们加载一个优化器，以将这些优化应用于我们的模型

>>> from optimum.onnxruntime import ORTOptimizer

>>> optimizer = ORTOptimizer.from_pretrained(ort_model)

>>> # Optimize the model
>>> optimizer.optimize(save_dir=save_directory, optimization_config=optimization_config)

就是这样 - 模型现在已优化并准备好进行推理！如您所见，每种情况下的过程都相似

通过 OptimizationConfig / QuantizationConfig 对象定义优化/量化策略
实例化 ORTQuantizer 或 ORTOptimizer 类
应用 quantize() 或 optimize() 方法
运行推理

查看 examples 目录以获取更复杂的使用示例。

祝您优化愉快 🤗！

< > 在 GitHub 上更新