将 🤗 Transformers 模型导出到 ONNX

🤗 Transformers 提供了一个 transformers.onnx 包，使您能够通过利用配置对象将模型检查点转换为 ONNX 图。

请参阅关于导出 🤗 Transformers 模型的指南以获取更多详细信息。

ONNX 配置

我们提供了三个抽象类，您应该从它们继承，具体取决于您希望导出的模型架构类型

基于编码器的模型从 OnnxConfig 继承
基于解码器的模型从 OnnxConfigWithPast 继承
编码器-解码器模型从 OnnxSeq2SeqConfigWithPast 继承

OnnxConfig

class transformers.onnx.OnnxConfig

( config: PretrainedConfig task: str = 'default' patching_specs: typing.List[transformers.onnx.config.PatchingSpec] = None )

用于 ONNX 可导出模型的基础类，描述了如何通过 ONNX 格式导出模型的元数据。

flatten_output_collection_property

< source >

( name: str field: typing.Iterable[typing.Any] ) → (Dict[str, Any])

参数

name — 嵌套结构的名称
field — 要展平的结构

(Dict[str, Any])

具有展平结构的输出和映射此新结构的键。

展平任何潜在的嵌套结构，并使用结构内元素的索引扩展字段名称。

from_model_config

< source >

( config: PretrainedConfig task: str = 'default' )

参数

config — 用于导出到 ONNX 的模型配置

实例化特定模型的 OnnxConfig

generate_dummy_inputs

< source >

( preprocessor: typing.Union[ForwardRef('PreTrainedTokenizerBase'), ForwardRef('FeatureExtractionMixin'), ForwardRef('ImageProcessingMixin')] batch_size: int = -1 seq_length: int = -1 num_choices: int = -1 is_pair: bool = False framework: typing.Optional[transformers.utils.generic.TensorType] = None num_channels: int = 3 image_width: int = 40 image_height: int = 40 sampling_rate: int = 22050 time_duration: float = 5.0 frequency: int = 220 tokenizer: PreTrainedTokenizerBase = None )

参数

preprocessor — (PreTrainedTokenizerBase, FeatureExtractionMixin, 或 ImageProcessingMixin): 与此模型配置关联的预处理器。
batch_size (int, 可选, 默认为 -1) — 用于导出模型的批次大小（-1 表示动态轴）。
num_choices (int, 可选, 默认为 -1) — 为多项选择任务提供的候选答案数量（-1 表示动态轴）。
seq_length (int, 可选, 默认为 -1) — 用于导出模型的序列长度（-1 表示动态轴）。
is_pair (bool, 可选, 默认为 False) — 指示输入是否为一对（句子 1，句子 2）
framework (TensorType, 可选, 默认为 None) — tokenizer 将为其生成张量的框架（PyTorch 或 TensorFlow）。
num_channels (int, 可选, 默认为 3) — 生成图像的通道数。
image_width (int, 可选, 默认为 40) — 生成图像的宽度。
image_height (int, 可选, 默认为 40) — 生成图像的高度。
sampling_rate (int, 可选默认为 22050) — 音频数据生成的采样率。
time_duration (float, 可选默认为 5.0) — 音频数据生成的总采样秒数。
frequency (int, optional defaults to 220) — 生成音频的所需自然频率 (int，可选，默认为 220)。

为特定框架生成提供给 ONNX 导出器的输入

generate_dummy_inputs_onnxruntime

< source >

( reference_model_inputs: typing.Mapping[str, typing.Any] ) → Mapping[str, Tensor]

参数

reference_model_inputs ([Mapping[str, Tensor]) — 模型的参考输入。

Mapping[str, Tensor]

保存提供给模型前向传播函数的 kwargs 的映射

使用参考模型输入为 ONNX Runtime 生成输入。覆盖此方法以使用序列到序列模型运行推理，这些模型的编码器和解码器导出为单独的 ONNX 文件。

use_external_data_format

< source >

( num_parameters: int )

参数

num_parameters — 模型上的参数数量

指示模型是否需要使用外部数据格式的标志

OnnxConfigWithPast

class transformers.onnx.OnnxConfigWithPast

< source >

( config: PretrainedConfig task: str = 'default' patching_specs: typing.List[transformers.onnx.config.PatchingSpec] = None use_past: bool = False )

fill_with_past_key_values_

< source >

( inputs_or_outputs: typing.Mapping[str, typing.Mapping[int, str]] direction: str inverted_values_shape: bool = False )

参数

inputs_or_outputs — 要填充的映射。
direction — “inputs” 或 “outputs” 之一，它指定 input_or_outputs 是输入映射还是输出映射，这对于轴命名非常重要。
inverted_values_shape — 如果为 True，则将值存储在动态轴 1 上，否则存储在轴 2 上。

使用 past_key_values 动态轴填充 input_or_outputs 映射，并考虑在内。

with_past

< source >

( config: PretrainedConfig task: str = 'default' )

参数

config — 底层模型的配置，用于导出到 ONNX 时使用

实例化一个 OnnxConfig，并将 use_past 属性设置为 True

OnnxSeq2SeqConfigWithPast

class transformers.onnx.OnnxSeq2SeqConfigWithPast

< source >

( config: PretrainedConfig task: str = 'default' patching_specs: typing.List[transformers.onnx.config.PatchingSpec] = None use_past: bool = False )

ONNX 特性

每个 ONNX 配置都与一组特性相关联，这些特性使您能够导出用于不同类型的拓扑或任务的模型。

FeaturesManager

class transformers.onnx.FeaturesManager

< source >

( )

check_supported_model_or_raise

< source >

( model: typing.Union[ForwardRef('PreTrainedModel'), ForwardRef('TFPreTrainedModel')] feature: str = 'default' )

参数

model — 要导出的模型。
feature — 要检查其是否可用的特性的名称。

检查模型是否具有请求的特性。

determine_framework

< source >

( model: str framework: str = None )

参数

model (str) — 要导出的模型的名称。
framework (str, optional, defaults to None) — 用于导出的框架。如果未提供，请参阅上面的优先级。

确定用于导出的框架。

优先级顺序如下

用户通过 framework 输入。
如果提供了本地检查点，则使用与检查点相同的框架。
环境中可用的框架，PyTorch 优先

get_config

< source >

( model_type: str feature: str ) → OnnxConfig

参数

model_type (str) — 要检索配置的模型类型。
feature (str) — 用于检索配置的特征。

OnnxConfig

组合的配置

获取模型类型和特征组合的 OnnxConfig。

get_model_class_for_feature

< source >

( feature: str framework: str = 'pt' )

参数

feature (str) — 所需的特征。
framework (str, 可选, 默认为 "pt") — 用于导出的框架。

尝试从特征名称检索 AutoModel 类。

get_model_from_feature

< source >

( feature: str model: str framework: str = None cache_dir: str = None )

参数

feature (str) — 所需的特征。
model (str) — 要导出的模型的名称。
framework (str, 可选, 默认为 None) — 用于导出的框架。如果未提供，请参阅 FeaturesManager.determine_framework 以了解优先级。

尝试从模型名称和要启用的特征检索模型。

get_supported_features_for_model_type

< source >

( model_type: str model_name: typing.Optional[str] = None )

参数

model_type (str) — 要检索支持特征的模型类型。
model_name (str, 可选) — 模型对象的名称属性，仅用于异常消息。

尝试从模型类型检索特征 -> OnnxConfig 构造函数映射。

< > 在 GitHub 上更新