Skip to content

Latest commit

 

History

History
63 lines (40 loc) · 2.55 KB

configuration.mdx

File metadata and controls

63 lines (40 loc) · 2.55 KB

Configuration classes for ONNX exports

Exporting a model to ONNX involves specifying:

  1. The input names.
  2. The output names.
  3. The dynamic axes. These refer to the input dimensions can be changed dynamically at runtime (e.g. a batch size or sequence length). All other axes will be treated as static, and hence fixed at runtime.
  4. Dummy inputs to trace the model. This is needed in PyTorch to record the computational graph and convert it to ONNX.

Since this data depends on the choice of model and task, we represent it in terms of configuration classes. Each configuration class is associated with a specific model architecture, and follows the naming convention ArchitectureNameOnnxConfig. For instance, the configuration which specifies the ONNX export of BERT models is BertOnnxConfig.

Since many architectures share similar properties for their ONNX configuration, 🤗 Optimum adopts a 3-level class hierarchy:

  1. Abstract and generic base classes. These handle all the fundamental features, while being agnostic to the modality (text, image, audio, etc).
  2. Middle-end classes. These are aware of the modality, but multiple can exist for the same modality depending on the inputs they support. They specify which input generators should be used for the dummy inputs, but remain model-agnostic.
  3. Model-specific classes like the BertOnnxConfig mentioned above. These are the ones actually used to export models.

Base classes

[[autodoc]] exporters.onnx.OnnxConfig - inputs - outputs - generate_dummy_inputs

[[autodoc]] exporters.onnx.OnnxConfigWithPast - add_past_key_values

[[autodoc]] exporters.onnx.OnnxSeq2SeqConfigWithPast

Middle-end classes

Text

[[autodoc]] exporters.onnx.config.TextEncoderOnnxConfig

[[autodoc]] exporters.onnx.config.TextDecoderOnnxConfig

[[autodoc]] exporters.onnx.config.TextSeq2SeqOnnxConfig

Vision

[[autodoc]] exporters.onnx.config.VisionOnnxConfig

Multi-modal

[[autodoc]] exporters.onnx.config.TextAndVisionOnnxConfig