Audio models support with optimum.exporters.onnx (#622)

* Audio models, first draft

* Hubert, Wav2Vec and Sew work

* Add support for the other models

* Almost done

* Add architecture names to doc

* Fixed issue with SEW-D and WavLM

* Add tiny tests

* Fix

* Add models to tiny

* Add Data2Vec audio

* Fix dummy input generator tests

* Update doc

* Fix tests

Co-authored-by: Michael Benayoun <michael@huggingface.co>

Author: michaelbenayoun

docs/source/exporters/onnx/package_reference/configuration.mdx

@@ -64,6 +64,7 @@ They specify which input generators should be used for the dummy inputs, but rem
 
 ## Supported architectures
 
+- Audio Spectrogram Transformer
 - Albert
 - Bart
 - Beit
@@ -78,6 +79,7 @@ They specify which input generators should be used for the dummy inputs, but rem
 - CodeGen
 - ConvBert
 - ConvNext
+- Data2VecAudio
 - Data2VecText
 - Data2VecVision
 - Deberta
@@ -91,6 +93,7 @@ They specify which input generators should be used for the dummy inputs, but rem
 - GPT-J
 - GPT-Neo
 - GroupVit
+- Hubert
 - IBert
 - LayoutLM
 - LayoutLM-v3
@@ -110,10 +113,17 @@ They specify which input generators should be used for the dummy inputs, but rem
 - Roberta
 - Roformer
 - Segformer
+- SEW
+- Speech2Text
 - SqueezeBert
 - Stable Diffusion
 - T5
+- UniSpeech
+- UniSpeech SAT
 - Vit
+- Wav2Vec2
+- Wav2Vec2 Conformer
+- WavLM
 - Whisper
 - XLM
 - XLM-Roberta

optimum/exporters/onnx/base.py

@@ -158,7 +158,11 @@ class OnnxConfig(ExportConfig, ABC):
         ),
         "sequence-classification": OrderedDict({"logits": {0: "batch_size"}}),
         "token-classification": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}),
-        "speech2seq-lm": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),
+        "speech2seq-lm": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}),
+        "audio-classification": OrderedDict({"logits": {0: "batch_size"}}),
+        "audio-frame-classification": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}),
+        "audio-ctc": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}),
+        "audio-xvector": OrderedDict({"logits": {0: "batch_size"}, "embeddings": {0: "batch_size"}}),
     }
 
     def __init__(

optimum/exporters/onnx/config.py

@@ -162,6 +162,10 @@ class AudioOnnxConfig(OnnxConfig):
 
     DUMMY_INPUT_GENERATOR_CLASSES = (DummyAudioInputGenerator,)
 
+    @property
+    def inputs(self) -> Mapping[str, Mapping[int, str]]:
+        return {"input_values": {0: "batch_size", 1: "sequence_length"}}
+
 
 class AudioToTextOnnxConfig(OnnxSeq2SeqConfigWithPast):
     DUMMY_INPUT_GENERATOR_CLASSES = (

optimum/exporters/onnx/model_configs.py

@@ -20,6 +20,7 @@
 
 from ...utils import (
     DEFAULT_DUMMY_SHAPES,
+    DummyAudioInputGenerator,
     DummyDecoderTextInputGenerator,
     DummyPastKeyValuesGenerator,
     DummySeq2SeqDecoderTextInputGenerator,
@@ -34,8 +35,9 @@
     NormalizedVisionConfig,
     logging,
 )
-from .base import ConfigBehavior, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast
+from .base import ConfigBehavior, OnnxConfig, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast
 from .config import (
+    AudioOnnxConfig,
     AudioToTextOnnxConfig,
     TextAndVisionOnnxConfig,
     TextDecoderOnnxConfig,
@@ -514,6 +516,18 @@ class SegformerOnnxConfig(YolosOnnxConfig):
     pass
 
 
+class MobileNetV1OnnxConfig(ViTOnnxConfig):
+    ATOL_FOR_VALIDATION = 1e-4
+
+    @property
+    def inputs(self) -> Mapping[str, Mapping[int, str]]:
+        return {"pixel_values": {0: "batch_size"}}
+
+
+class MobileNetV2OnnxConfig(MobileNetV1OnnxConfig):
+    pass
+
+
 class CLIPNormalizedConfig(NormalizedTextAndVisionConfig):
     TEXT_CONFIG = "text_config"
     VISION_CONFIG = "vision_config"
@@ -693,11 +707,9 @@ class Data2VecVisionOnnxConfig(ViTOnnxConfig):
     pass
 
 
-# TODO: add support when audio models are supported.
-class Data2VecAudioOnnxConfig(ViTOnnxConfig):
-    @property
-    def inputs(self):
-        raise NotImplementedError
+class Data2VecAudioOnnxConfig(AudioOnnxConfig):
+    NORMALIZED_CONFIG_CLASS = NormalizedConfig
+    ATOL_FOR_VALIDATION = 1e-4
 
 
 class PerceiverDummyInputGenerator(DummyVisionInputGenerator):
@@ -751,20 +763,94 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs):
         return dummy_inputs
 
 
-class WhisperOnnxConfig(AudioToTextOnnxConfig):
-    NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig
-    ATOL_FOR_VALIDATION = 1e-3
+class HubertOnnxConfig(AudioOnnxConfig):
+    NORMALIZED_CONFIG_CLASS = NormalizedConfig
 
 
-class MobileNetV1OnnxConfig(VisionOnnxConfig):
-    NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig
-    MIN_TORCH_VERSION = version.parse("1.11")
+class Wav2Vec2OnnxConfig(HubertOnnxConfig):
+    pass
+
+
+class Wav2Vec2ConformerOnnxConfig(HubertOnnxConfig):
+    pass
+
+
+class SEWOnnxConfig(HubertOnnxConfig):
+    pass
+
+
+class SEWDOnnxConfig(HubertOnnxConfig):
+    DEFAULT_ONNX_OPSET = 12
+
+
+class UniSpeechOnnxConfig(HubertOnnxConfig):
+    pass
+
+
+class UniSpeechSATOnnxConfig(HubertOnnxConfig):
+    pass
+
+
+class WavLMOnnxConfig(HubertOnnxConfig):
+    DEFAULT_ONNX_OPSET = 12
+
+
+class ASTDummyAudioInputGenerator(DummyAudioInputGenerator):
+    def generate(self, input_name: str, framework: str = "pt"):
+        shape = [self.batch_size, self.normalized_config.max_length, self.normalized_config.num_mel_bins]
+        if input_name == "input_values":
+            return self.random_float_tensor(shape, min_value=-1, max_value=1, framework=framework)
+        return super().generate(input_name, framework=framework)
+
+
+class ASTOnnxConfig(OnnxConfig):
+    NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args(
+        num_mel_bins="num_mel_bins", max_length="max_length", allow_new=True
+    )
+    DUMMY_INPUT_GENERATOR_CLASSES = (ASTDummyAudioInputGenerator,)
     ATOL_FOR_VALIDATION = 1e-4
 
     @property
     def inputs(self) -> Mapping[str, Mapping[int, str]]:
-        return {"pixel_values": {0: "batch"}}
+        return {"input_values": {0: "batch_size"}}
 
 
-class MobileNetV2OnnxConfig(MobileNetV1OnnxConfig):
-    pass
+# TODO: currently disabled because an operator seems not supported by ONNX.
+# class MCTCTDummyAudioInputGenerator(DummyAudioInputGenerator):
+#     def generate(self, input_name: str, framework: str = "pt"):
+#         shape = [self.batch_size, self.sequence_length, self.normalized_config.input_features_per_channel]
+#         if input_name == "input_features":
+#             return self.random_float_tensor(shape, min_value=-1, max_value=1, framework=framework)
+#         return super().generate(input_name, framework=framework)
+#
+#
+# class MCTCTOnnxConfig(OnnxConfig):
+#     NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args(input_features_per_channel="input_feat_per_channel", allow_new=True)
+#     DUMMY_INPUT_GENERATOR_CLASSES = (MCTCTDummyAudioInputGenerator,)
+#     DEFAULT_ONNX_OPSET = 13
+#
+#     @property
+#     def inputs(self) -> Mapping[str, Mapping[int, str]]:
+#         return {"input_features": {0: "batch_size", 1: "sequence_classification"}}
+
+
+class WhisperOnnxConfig(AudioToTextOnnxConfig):
+    NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig
+    ATOL_FOR_VALIDATION = 1e-3
+
+
+class Speech2TextDummyAudioInputGenerator(DummyAudioInputGenerator):
+    def generate(self, input_name: str, framework: str = "pt"):
+        shape = [self.batch_size, self.sequence_length, self.normalized_config.input_features_per_channel]
+        if input_name == "input_features":
+            return self.random_float_tensor(shape, min_value=-1, max_value=1, framework=framework)
+        return super().generate(input_name, framework=framework)
+
+
+class Speech2TextOnnxConfig(AudioToTextOnnxConfig):
+    NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args(
+        input_features_per_channel="input_feat_per_channel", allow_new=True
+    )
+    DUMMY_INPUT_GENERATOR_CLASSES = (
+        Speech2TextDummyAudioInputGenerator,
+    ) + AudioToTextOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES[1:]

optimum/exporters/tasks.py

@@ -100,6 +100,10 @@ class TasksManager:
             "masked-im": "AutoModelForMaskedImageModeling",
             "semantic-segmentation": "AutoModelForSemanticSegmentation",
             "speech2seq-lm": "AutoModelForSpeechSeq2Seq",
+            "audio-classification": "AutoModelForAudioClassification",
+            "audio-frame-classification": "AutoModelForAudioFrameClassification",
+            "audio-ctc": "AutoModelForCTC",
+            "audio-xvector": "AutoModelForAudioXVector",
             "stable-diffusion": "StableDiffusionPipeline",
         }
     if is_tf_available():
@@ -130,11 +134,20 @@ class TasksManager:
         "masked-im": "transformers",
         "semantic-segmentation": "transformers",
         "speech2seq-lm": "transformers",
+        "audio-ctc": "transformers",
+        "audio-classification": "transformers",
+        "audio-frame-classification": "transformers",
+        "audio-xvector": "transformers",
         "stable-diffusion": "diffusers",
     }
 
     # Set of model topologies we support associated to the tasks supported by each topology and the factory
     _SUPPORTED_MODEL_TYPE = {
+        "audio-spectrogram-transformer": supported_tasks_mapping(
+            "default",
+            "audio-classification",
+            onnx="ASTOnnxConfig",
+        ),
         "albert": supported_tasks_mapping(
             "default",
             "masked-lm",
@@ -273,6 +286,14 @@ class TasksManager:
             # "semantic-segmentation",
             onnx="Data2VecVisionOnnxConfig",
         ),
+        "data2vec-audio": supported_tasks_mapping(
+            "default",
+            "audio-ctc",
+            "audio-classification",
+            "audio-frame-classification",
+            "audio-xvector",
+            onnx="Data2VecAudioOnnxConfig",
+        ),
         "deberta": supported_tasks_mapping(
             "default",
             "masked-lm",
@@ -356,6 +377,12 @@ class TasksManager:
             "default",
             onnx="GroupViTOnnxConfig",
         ),
+        "hubert": supported_tasks_mapping(
+            "default",
+            "audio-ctc",
+            "audio-classification",
+            onnx="HubertOnnxConfig",
+        ),
         "ibert": supported_tasks_mapping(
             "default",
             "masked-lm",
@@ -423,6 +450,12 @@ class TasksManager:
             "question-answering",
             onnx="MBartOnnxConfig",
         ),
+        # TODO: enable once the missing operator is supported.
+        # "mctct": supported_tasks_mapping(
+        #     "default",
+        #     "audio-ctc",
+        #     onnx="MCTCTOnnxConfig",
+        # ),
         "mobilebert": supported_tasks_mapping(
             "default",
             "masked-lm",
@@ -521,6 +554,25 @@ class TasksManager:
             "semantic-segmentation",
             onnx="SegformerOnnxConfig",
         ),
+        "sew": supported_tasks_mapping(
+            "default",
+            "audio-ctc",
+            "audio-classification",
+            onnx="SEWOnnxConfig",
+        ),
+        "sew-d": supported_tasks_mapping(
+            "default",
+            "audio-ctc",
+            "audio-classification",
+            onnx="SEWDOnnxConfig",
+        ),
+        "speech-to-text": supported_tasks_mapping(
+            "default",
+            "default-with-past",
+            "speech2seq-lm",
+            "speech2seq-lm-with-past",
+            onnx="Speech2TextOnnxConfig",
+        ),
         "squeezebert": supported_tasks_mapping(
             "default",
             "masked-lm",
@@ -530,6 +582,12 @@ class TasksManager:
             "question-answering",
             onnx="SqueezeBertOnnxConfig",
         ),
+        "swin": supported_tasks_mapping(
+            "default",
+            "image-classification",
+            "masked-im",
+            onnx="SwinOnnxConfig",
+        ),
         "t5": supported_tasks_mapping(
             "default",
             "default-with-past",
@@ -541,11 +599,49 @@ class TasksManager:
             "semantic-segmentation",
             onnx="UNetOnnxConfig",
         ),
+        "unispeech": supported_tasks_mapping(
+            "default",
+            "audio-ctc",
+            "audio-classification",
+            onnx="UniSpeechOnnxConfig",
+        ),
+        "unispeech-sat": supported_tasks_mapping(
+            "default",
+            "audio-ctc",
+            "audio-classification",
+            "audio-frame-classification",
+            "audio-xvector",
+            onnx="UniSpeechSATOnnxConfig",
+        ),
         "vae": supported_tasks_mapping(
             "semantic-segmentation",
             onnx="VaeOnnxConfig",
         ),
         "vit": supported_tasks_mapping("default", "image-classification", "masked-im", onnx="ViTOnnxConfig"),
+        "wavlm": supported_tasks_mapping(
+            "default",
+            "audio-ctc",
+            "audio-classification",
+            "audio-frame-classification",
+            "audio-xvector",
+            onnx="WavLMOnnxConfig",
+        ),
+        "wav2vec2": supported_tasks_mapping(
+            "default",
+            "audio-ctc",
+            "audio-classification",
+            "audio-frame-classification",
+            "audio-xvector",
+            onnx="Wav2Vec2OnnxConfig",
+        ),
+        "wav2vec2-conformer": supported_tasks_mapping(
+            "default",
+            "audio-ctc",
+            "audio-classification",
+            "audio-frame-classification",
+            "audio-xvector",
+            onnx="Wav2Vec2ConformerOnnxConfig",
+        ),
         "whisper": supported_tasks_mapping(
             "default",
             "default-with-past",
@@ -580,12 +676,6 @@ class TasksManager:
             "object-detection",
             onnx="YolosOnnxConfig",
         ),
-        "swin": supported_tasks_mapping(
-            "default",
-            "image-classification",
-            "masked-im",
-            onnx="SwinOnnxConfig",
-        ),
     }
     _UNSUPPORTED_CLI_MODEL_TYPE = {"unet", "vae", "clip-text-model"}
     _SUPPORTED_CLI_MODEL_TYPE = set(_SUPPORTED_MODEL_TYPE.keys()) - _UNSUPPORTED_CLI_MODEL_TYPE