Hugging Face's logo

Transformers documentation

Exporting 🤗 Transformers models to ONNX

You are viewing v4.22.2 version. A newer version v4.46.2 is available.
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Exporting 🤗 Transformers models to ONNX

🤗 Transformers provides a transformers.onnx package that enables you to convert model checkpoints to an ONNX graph by leveraging configuration objects.

See the guide on exporting 🤗 Transformers models for more details.

ONNX Configurations

We provide three abstract classes that you should inherit from, depending on the type of model architecture you wish to export:

OnnxConfig

class transformers.onnx.OnnxConfig

< >

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

Base class for ONNX exportable model describing metadata on how to export the model through the ONNX format.

flatten_output_collection_property

< >

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

Returns

(Dict[str, Any])

Outputs with flattened structure and key mapping this new structure.

Flatten any potential nested structure expanding the name of the field with the index of the element within the structure.

from_model_config

< >

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

Instantiate a OnnxConfig for a specific model

generate_dummy_inputs

< >

( preprocessor: typing.Union[ForwardRef('PreTrainedTokenizerBase'), ForwardRef('FeatureExtractionMixin')] 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 tokenizer: PreTrainedTokenizerBase = None )

Parameters

  • batch_size (int, optional, defaults to -1) — The batch size to export the model for (-1 means dynamic axis).
  • num_choices (int, optional, defaults to -1) — The number of candidate answers provided for multiple choice task (-1 means dynamic axis).
  • seq_length (int, optional, defaults to -1) — The sequence length to export the model for (-1 means dynamic axis).
  • is_pair (bool, optional, defaults to False) — Indicate if the input is a pair (sentence 1, sentence 2)
  • framework (TensorType, optional, defaults to None) — The framework (PyTorch or TensorFlow) that the tokenizer will generate tensors for.
  • num_channels (int, optional, defaults to 3) — The number of channels of the generated images.
  • image_width (int, optional, defaults to 40) — The width of the generated images.
  • image_height (int, optional, defaults to 40) — The height of the generated images.

Generate inputs to provide to the ONNX exporter for the specific framework

use_external_data_format

< >

( num_parameters: int )

Flag indicating if the model requires using external data format

OnnxConfigWithPast

class transformers.onnx.OnnxConfigWithPast

< >

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

fill_with_past_key_values_

< >

( inputs_or_outputs: typing.Mapping[str, typing.Mapping[int, str]] direction: str )

Fill the input_or_outputs mapping with past_key_values dynamic axes considering.

with_past

< >

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

Instantiate a OnnxConfig with use_past attribute set to True

OnnxSeq2SeqConfigWithPast

class transformers.onnx.OnnxSeq2SeqConfigWithPast

< >

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

ONNX Features

Each ONNX configuration is associated with a set of features that enable you to export models for different types of topologies or tasks.

FeaturesManager

class transformers.onnx.FeaturesManager

< >

( )

check_supported_model_or_raise

< >

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

Check whether or not the model has the requested features.

determine_framework

< >

( model: str framework: str = None )

Parameters

  • model (str) — The name of the model to export.
  • framework (str, optional, defaults to None) — The framework to use for the export. See above for priority if none provided.

Determines the framework to use for the export.

The priority is in the following order:

  1. User input via framework.
  2. If local checkpoint is provided, use the same framework as the checkpoint.
  3. Available framework in environment, with priority given to PyTorch

get_config

< >

( model_type: str feature: str ) → OnnxConfig

Parameters

  • model_type (str) — The model type to retrieve the config for.
  • feature (str) — The feature to retrieve the config for.

Returns

OnnxConfig

config for the combination

Gets the OnnxConfig for a model_type and feature combination.

get_model_class_for_feature

< >

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

Parameters

  • feature (str) — The feature required.
  • framework (str, optional, defaults to "pt") — The framework to use for the export.

Attempts to retrieve an AutoModel class from a feature name.

get_model_from_feature

< >

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

Parameters

  • feature (str) — The feature required.
  • model (str) — The name of the model to export.
  • framework (str, optional, defaults to None) — The framework to use for the export. See FeaturesManager.determine_framework for the priority should none be provided.

Attempts to retrieve a model from a model’s name and the feature to be enabled.

get_supported_features_for_model_type

< >

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

Parameters

  • model_type (str) — The model type to retrieve the supported features for.
  • model_name (str, optional) — The name attribute of the model object, only used for the exception message.

Tries to retrieve the feature -> OnnxConfig constructor map from the model type.

←Text Generation Optimization→