optimum-onnx 文档

优化

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

优化

🤗 Optimum 提供了一个 optimum.onnxruntime 包,使您能够使用 ONNX Runtime 模型优化工具,对托管在 🤗 hub 上的许多模型应用图优化。

在 ONNX 导出期间优化模型

ONNX 模型可以在 ONNX 导出期间使用 Optimum CLI 直接进行优化,通过在 CLI 中传递参数 --optimize {O1,O2,O3,O4},例如

optimum-cli export onnx --model gpt2 --optimize O3 gpt2_onnx/

优化级别如下:

  • O1:基础通用优化。
  • O2:基础和扩展通用优化,以及针对 transformers 的特定融合(fusions)。
  • O3:与 O2 相同,但包含 GELU 近似。
  • O4:与 O3 相同,但包含混合精度 (fp16, 仅限 GPU, 需要 --device cuda)。

使用 ORTOptimizer 通过编程方式优化模型

ONNX 模型可以使用 ORTOptimizer 进行优化。该类可以使用 from_pretrained() 方法初始化,该方法支持不同的检查点格式。

  1. 使用已经初始化的 ORTModel 类。
>>> from optimum.onnxruntime import ORTOptimizer, ORTModelForSequenceClassification

# Loading ONNX Model from the Hub
>>> model = ORTModelForSequenceClassification.from_pretrained(
...     "optimum/distilbert-base-uncased-finetuned-sst-2-english"
... )

# Create an optimizer from an ORTModelForXXX
>>> optimizer = ORTOptimizer.from_pretrained(model)
  1. 使用来自目录的本地 ONNX 模型。
>>> from optimum.onnxruntime import ORTOptimizer

# This assumes a model.onnx exists in path/to/model
>>> optimizer = ORTOptimizer.from_pretrained("path/to/model")

优化配置

OptimizationConfig 类允许指定 ORTOptimizer 应该如何执行优化。

在优化配置中,有 4 个可能的优化级别:

  • optimization_level=0:禁用所有优化
  • optimization_level=1:启用基础优化,例如常量折叠或冗余节点消除
  • optimization_level=2:启用扩展图优化,例如节点融合
  • optimization_level=99:启用数据布局优化

选择某个级别将启用该级别的优化以及所有前置级别的优化。更多信息请参阅 此处

enable_transformers_specific_optimizations=True 意味着除了上述 ONNX Runtime 优化外,还将执行 transformers 特定的图融合和近似。以下是您可以启用的可能优化列表:

  • 通过 disable_gelu_fusion=False 启用 Gelu 融合,
  • 通过 disable_layer_norm_fusion=False 启用层归一化(Layer Normalization)融合,
  • 通过 disable_attention_fusion=False 启用注意力(Attention)融合,
  • 通过 disable_skip_layer_norm_fusion=False 启用 SkipLayerNormalization 融合,
  • 通过 disable_bias_skip_layer_norm_fusion=False 启用 Add Bias 和 SkipLayerNormalization 融合,
  • 通过 disable_bias_gelu_fusion=False 启用 Add Bias 和 Gelu / FastGelu 融合,
  • 通过 enable_gelu_approximation=True 启用 Gelu 近似。

注意力融合是为类 BERT 架构(如 BERT, RoBERTa, VIT 等)的右侧填充(right-side padding)和生成式模型(类 GPT)的左侧填充(left-side padding)设计的。如果您不遵循此约定,请设置 use_raw_attention_mask=True 以避免潜在的精度问题,但会牺牲性能。

虽然 OptimizationConfig 让您能够完全控制如何进行优化,但很难知道应该启用/禁用什么。相反,您可以使用 AutoOptimizationConfig,它提供了四个常用的优化级别:

  • O1:基础通用优化。
  • O2:基础和扩展通用优化,以及针对 transformers 的特定融合(fusions)。
  • O3:与 O2 相同,但包含 GELU 近似。
  • O4:与 O3 相同,但包含混合精度 (fp16, 仅限 GPU)。

示例:加载一个 O2 OptimizationConfig

>>> from optimum.onnxruntime import AutoOptimizationConfig
>>> optimization_config = AutoOptimizationConfig.O2()

您还可以指定 O2 配置中未定义的自定义参数,例如:

>>> from optimum.onnxruntime import AutoOptimizationConfig
>>> optimization_config = AutoOptimizationConfig.O2(disable_embed_layer_norm_fusion=False)

优化示例

下面您将看到一个关于如何优化 distilbert-base-uncased-finetuned-sst-2-english 的简单端到端示例。

>>> from optimum.onnxruntime import (
...     AutoOptimizationConfig, ORTOptimizer, ORTModelForSequenceClassification
... )

>>> model_id = "distilbert-base-uncased-finetuned-sst-2-english"
>>> save_dir = "distilbert_optimized"

>>> # Load a PyTorch model and export it to the ONNX format
>>> model = ORTModelForSequenceClassification.from_pretrained(model_id, export=True)

>>> # Create the optimizer
>>> optimizer = ORTOptimizer.from_pretrained(model)

>>> # Define the optimization strategy by creating the appropriate configuration
>>> optimization_config = AutoOptimizationConfig.O2()

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

下面您将看到一个关于如何优化 Seq2Seq 模型 sshleifer/distilbart-cnn-12-6 的简单端到端示例。

>>> from transformers import AutoTokenizer
>>> from optimum.onnxruntime import  OptimizationConfig, ORTOptimizer, ORTModelForSeq2SeqLM

>>> model_id = "sshleifer/distilbart-cnn-12-6"
>>> save_dir = "distilbart_optimized"

>>> # Load a PyTorch model and export it to the ONNX format
>>> model = ORTModelForSeq2SeqLM.from_pretrained(model_id, export=True)

>>> # Create the optimizer
>>> optimizer = ORTOptimizer.from_pretrained(model)

>>> # Define the optimization strategy by creating the appropriate configuration
>>> optimization_config = OptimizationConfig(
...     optimization_level=2,
...     enable_transformers_specific_optimizations=True,
...     optimize_for_gpu=False,
... )

>>> # Optimize the model
>>> optimizer.optimize(save_dir=save_dir, optimization_config=optimization_config)
>>> tokenizer = AutoTokenizer.from_pretrained(model_id)
>>> optimized_model = ORTModelForSeq2SeqLM.from_pretrained(save_dir)
>>> tokens = tokenizer("This is a sample input", return_tensors="pt")
>>> outputs = optimized_model.generate(**tokens)

使用 Optimum CLI 优化模型

Optimum ONNX Runtime 优化工具可以直接通过 Optimum 命令行界面使用

optimum-cli onnxruntime optimize --help
usage: optimum-cli <command> [<args>] onnxruntime optimize [-h] --onnx_model ONNX_MODEL -o OUTPUT (-O1 | -O2 | -O3 | -O4 | -c CONFIG)

options:
  -h, --help            show this help message and exit
  -O1                   Basic general optimizations (see: https://huggingface.co/docs/optimum/onnxruntime/usage_guides/optimization for more details).
  -O2                   Basic and extended general optimizations, transformers-specific fusions (see: https://huggingface.co/docs/optimum/onnxruntime/usage_guides/optimization for more
                        details).
  -O3                   Same as O2 with Gelu approximation (see: https://huggingface.co/docs/optimum/onnxruntime/usage_guides/optimization for more details).
  -O4                   Same as O3 with mixed precision (see: https://huggingface.co/docs/optimum/onnxruntime/usage_guides/optimization for more details).
  -c CONFIG, --config CONFIG
                        `ORTConfig` file to use to optimize the model.

Required arguments:
  --onnx_model ONNX_MODEL
                        Path to the repository where the ONNX models to optimize are located.
  -o OUTPUT, --output OUTPUT
                        Path to the directory where to store generated ONNX model.

优化 ONNX 模型可以按如下方式进行

 optimum-cli onnxruntime optimize --onnx_model onnx_model_location/ -O1 -o optimized_model/

这将使用基础通用优化对 onnx_model_location 中的所有 ONNX 文件进行优化。

© . This site is unofficial and not affiliated with Hugging Face, Inc.