Transformers documentation

Exporting 🤗 Transformers models to ONNX

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' )

Returns

OnnxConfig for this model

Instantiate a OnnxConfig for a specific model

generate_dummy_inputs < >

( tokenizer: PreTrainedTokenizer batch_size: int = -1 seq_length: int = -1 is_pair: bool = False framework: typing.Optional[transformers.file_utils.TensorType] = None )

Returns

Mapping[str, Tensor] holding the kwargs to provide to the model’s forward function

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

use_external_data_format < >

( num_parameters: int )

Returns

True if model.num_parameters() * size_of(float32) >= 2Gb False otherwise

Flag indicating if the model requires using external data format

OnnxConfigWithPast

fill_with_past_key_values_ < >

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

Fill the input_or_ouputs mapping with past_key_values dynamic axes considering.

with_past < >

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

Returns

OnnxConfig with .use_past = True

Instantiate a OnnxConfig with use_past attribute set to True

OnnxSeq2SeqConfigWithPast

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

check_supported_model_or_raise < >

( model: PreTrainedModel feature: str = 'default' )

Returns

(str) The type of the model (OnnxConfig) The OnnxConfig instance holding the model export properties.

Check whether or not the model has the requested features.

get_model_class_for_feature < >

( feature: str )

Returns

The AutoModel class corresponding to the feature.

Attempt to retrieve an AutoModel class from a feature name.

get_model_from_feature < >

( feature: str model: str )

Returns

The instance of the model.

Attempt 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 )

Returns

The dictionary mapping each feature to a corresponding OnnxConfig constructor.

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