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:
- Encoder-based models inherit from OnnxConfig
- Decoder-based models inherit from OnnxConfigWithPast
- Encoder-decoder models inherit from OnnxSeq2SeqConfigWithPast
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.
( 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.
( config: PretrainedConfig task: str = 'default' )
Returns
OnnxConfig for this model
Instantiate a OnnxConfig for a specific model
( 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
( 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
( 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.
( 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
( 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.
( feature: str )
Returns
The AutoModel class corresponding to the feature.
Attempt to retrieve an AutoModel class from a feature name.
Attempt to retrieve a model from a model’s name and the feature to be enabled.
( 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.