Add ONNX export for ViT

docs/source/serialization.mdx

@@ -59,6 +59,7 @@ Ready-made configurations include the following architectures:
 - OpenAI GPT-2
 - RoBERTa
 - T5
+- ViT
 - XLM-RoBERTa
 - XLM-RoBERTa-XL
 

src/transformers/models/vit/__init__.py

@@ -21,7 +21,7 @@
 
 
 _import_structure = {
-    "configuration_vit": ["VIT_PRETRAINED_CONFIG_ARCHIVE_MAP", "ViTConfig"],
+    "configuration_vit": ["VIT_PRETRAINED_CONFIG_ARCHIVE_MAP", "ViTConfig", "ViTOnnxConfig"],
 }
 
 if is_vision_available():
@@ -50,7 +50,7 @@
     ]
 
 if TYPE_CHECKING:
-    from .configuration_vit import VIT_PRETRAINED_CONFIG_ARCHIVE_MAP, ViTConfig
+    from .configuration_vit import VIT_PRETRAINED_CONFIG_ARCHIVE_MAP, ViTConfig, ViTOnnxConfig
 
     if is_vision_available():
         from .feature_extraction_vit import ViTFeatureExtractor

src/transformers/models/vit/configuration_vit.py

@@ -14,7 +14,11 @@
 # limitations under the License.
 """ ViT model configuration"""
 
+from collections import OrderedDict
+from typing import Mapping
+
 from ...configuration_utils import PretrainedConfig
+from ...onnx import OnnxConfig
 from ...utils import logging
 
 
@@ -117,3 +121,13 @@ def __init__(
         self.patch_size = patch_size
         self.num_channels = num_channels
         self.qkv_bias = qkv_bias
+
+
+class ViTOnnxConfig(OnnxConfig):
+    @property
+    def inputs(self) -> Mapping[str, Mapping[int, str]]:
+        return OrderedDict(
+            [
+                ("pixel_values", {0: "batch", 1: "sequence"}),
+            ]
+        )

src/transformers/onnx/__main__.py

@@ -15,7 +15,7 @@
 from argparse import ArgumentParser
 from pathlib import Path
 
-from transformers.models.auto import AutoTokenizer
+from transformers.models.auto import AutoFeatureExtractor, AutoTokenizer
 
 from ..utils import logging
 from .convert import export, validate_model_outputs
@@ -47,10 +47,16 @@ def main():
         args.output.parent.mkdir(parents=True)
 
     # Allocate the model
-    tokenizer = AutoTokenizer.from_pretrained(args.model)
     model = FeaturesManager.get_model_from_feature(args.feature, args.model)
     model_kind, model_onnx_config = FeaturesManager.check_supported_model_or_raise(model, feature=args.feature)
     onnx_config = model_onnx_config(model.config)
+    # Check the modality of the inputs and instantiate the appropriate preprocessor
+    if model.main_input_name == "input_ids":
+        preprocessor = AutoTokenizer.from_pretrained(args.model)
+    elif model.main_input_name == "pixel_values":
+        preprocessor = AutoFeatureExtractor.from_pretrained(args.model)
+    else:
+        raise ValueError(f"Unsupported model input name: {model.main_input_name}")
 
     # Ensure the requested opset is sufficient
     if args.opset is None:
@@ -62,12 +68,18 @@ def main():
             f"At least  {onnx_config.default_onnx_opset} is required."
         )
 
-    onnx_inputs, onnx_outputs = export(tokenizer, model, onnx_config, args.opset, args.output)
+    onnx_inputs, onnx_outputs = export(
+        preprocessor,
+        model,
+        onnx_config,
+        args.opset,
+        args.output,
+    )
 
     if args.atol is None:
         args.atol = onnx_config.atol_for_validation
 
-    validate_model_outputs(onnx_config, tokenizer, model, args.output, onnx_outputs, args.atol)
+    validate_model_outputs(onnx_config, preprocessor, model, args.output, onnx_outputs, args.atol)
     logger.info(f"All good, model saved at: {args.output.as_posix()}")
 
 

src/transformers/onnx/config.py

@@ -15,11 +15,24 @@
 import dataclasses
 from abc import ABC, abstractmethod
 from collections import OrderedDict
-from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple
+from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple, Union
+
+import requests
+from transformers import (
+    PretrainedConfig,
+    PreTrainedTokenizer,
+    PreTrainedTokenizerFast,
+    TensorType,
+    is_torch_available,
+    is_vision_available,
+)
+
+from ..feature_extraction_utils import FeatureExtractionMixin
+from .utils import ParameterFormat, compute_effective_axis_dimension, compute_serialized_parameters_size
 
-from transformers import PretrainedConfig, PreTrainedTokenizer, TensorType, is_torch_available
 
-from .utils import ParameterFormat, compute_effective_axis_dimension, compute_serialized_parameters_size
+if is_vision_available():
+    from PIL import Image
 
 
 DEFAULT_ONNX_OPSET = 11
@@ -71,6 +84,7 @@ class OnnxConfig(ABC):
                 "end_logits": {0: "batch", 1: "sequence"},
             }
         ),
+        "image-classification": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),
     }
 
     def __init__(self, config: PretrainedConfig, task: str = "default", patching_specs: List[PatchingSpec] = None):
@@ -197,40 +211,60 @@ def use_external_data_format(num_parameters: int) -> bool:
 
     def generate_dummy_inputs(
         self,
-        tokenizer: PreTrainedTokenizer,
+        preprocessor: Union[PreTrainedTokenizer, FeatureExtractionMixin],
         batch_size: int = -1,
         seq_length: int = -1,
         is_pair: bool = False,
         framework: Optional[TensorType] = None,
+        tokenizer: PreTrainedTokenizer = None,
     ) -> Mapping[str, Any]:
         """
         Generate inputs to provide to the ONNX exporter for the specific framework
 
         Args:
-            tokenizer: The tokenizer associated with this model configuration
-            batch_size: The batch size (int) to export the model for (-1 means dynamic axis)
-            seq_length: The sequence length (int) to export the model for (-1 means dynamic axis)
-            is_pair: Indicate if the input is a pair (sentence 1, sentence 2)
-            framework: The framework (optional) the tokenizer will generate tensor for
+            preprocessor: ([`PreTrainedTokenizer`] or [`FeatureExtractionMixin`]):
+                The preprocessor associated with this model configuration.
+            batch_size (`int`):
+                The batch size (int) to export the model for (-1 means dynamic axis)
+            seq_length (`int`):
+                The sequence length (int) to export the model for (-1 means dynamic axis)
+            is_pair (`bool`):
+                Indicate if the input is a pair (sentence 1, sentence 2)
+            framework (`TensorType`):
+                The framework (optional) the tokenizer will generate tensor for
+            tokenizer ([`PreTrainedTokenizer`]):
+                The tokenizer associated with this model configuration
 
         Returns:
             Mapping[str, Tensor] holding the kwargs to provide to the model's forward function
         """
+        if isinstance(preprocessor, PreTrainedTokenizer) and tokenizer:
+            raise ValueError("You cannot provide both a tokenizer and a preprocessor to generate dummy inputs.")
+        if isinstance(preprocessor, PreTrainedTokenizer) or isinstance(preprocessor, PreTrainedTokenizerFast):
+            if tokenizer:
+                preprocessor = tokenizer
+            # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
+            batch_size = compute_effective_axis_dimension(
+                batch_size, fixed_dimension=OnnxConfig.DEFAULT_FIXED_BATCH, num_token_to_add=0
+            )
 
-        # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
-        batch_size = compute_effective_axis_dimension(
-            batch_size, fixed_dimension=OnnxConfig.DEFAULT_FIXED_BATCH, num_token_to_add=0
-        )
-
-        # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
-        token_to_add = tokenizer.num_special_tokens_to_add(is_pair)
-        seq_length = compute_effective_axis_dimension(
-            seq_length, fixed_dimension=OnnxConfig.DEFAULT_FIXED_SEQUENCE, num_token_to_add=token_to_add
-        )
+            # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
+            token_to_add = preprocessor.num_special_tokens_to_add(is_pair)
+            seq_length = compute_effective_axis_dimension(
+                seq_length, fixed_dimension=OnnxConfig.DEFAULT_FIXED_SEQUENCE, num_token_to_add=token_to_add
+            )
 
-        # Generate dummy inputs according to compute batch and sequence
-        dummy_input = [" ".join([tokenizer.unk_token]) * seq_length] * batch_size
-        return dict(tokenizer(dummy_input, return_tensors=framework))
+            # Generate dummy inputs according to compute batch and sequence
+            dummy_input = [" ".join([preprocessor.unk_token]) * seq_length] * batch_size
+            return dict(preprocessor(dummy_input, return_tensors=framework))
+        elif isinstance(preprocessor, FeatureExtractionMixin) and is_vision_available():
+            url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+            image = Image.open(requests.get(url, stream=True).raw)
+            return dict(preprocessor(images=image, return_tensors=framework))
+        else:
+            raise ValueError(
+                "Unable to generate dummy inputs for the model. Please provide a tokenizer or a preprocessor."
+            )
 
     def patch_ops(self):
         for spec in self._patching_specs:

src/transformers/onnx/convert.py

@@ -25,6 +25,8 @@
 from transformers.onnx.config import OnnxConfig
 from transformers.utils import logging
 
+from ..feature_extraction_utils import FeatureExtractionMixin
+
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
@@ -63,26 +65,29 @@ def check_onnxruntime_requirements(minimum_version: Version):
 
 
 def export_pytorch(
-    tokenizer: PreTrainedTokenizer,
+    preprocessor: Union[PreTrainedTokenizer, FeatureExtractionMixin],
     model: PreTrainedModel,
     config: OnnxConfig,
     opset: int,
     output: Path,
+    tokenizer: PreTrainedTokenizer = None,
 ) -> Tuple[List[str], List[str]]:
     """
     Export a PyTorch model to an ONNX Intermediate Representation (IR)
 
     Args:
-        tokenizer ([`PreTrainedTokenizer`]):
-            The tokenizer used for encoding the data.
-        model ([`PreTrainedModel`]):
+        preprocessor: ([`PreTrainedTokenizer`] or [`FeatureExtractionMixin`]):
+            The preprocessor used for encoding the data.
+        model ([`PreTrainedModel`] or [`TFPreTrainedModel`]):
             The model to export.
         config ([`~onnx.config.OnnxConfig`]):
             The ONNX configuration associated with the exported model.
         opset (`int`):
             The version of the ONNX operator set to use.
         output (`Path`):
             Directory to store the exported ONNX model.
+        tokenizer ([`PreTrainedTokenizer`]):
+            The tokenizer used for encoding the data.
 
     Returns:
         `Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from
@@ -106,7 +111,9 @@ def export_pytorch(
 
             # Ensure inputs match
             # TODO: Check when exporting QA we provide "is_pair=True"
-            model_inputs = config.generate_dummy_inputs(tokenizer, framework=TensorType.PYTORCH)
+            model_inputs = config.generate_dummy_inputs(
+                preprocessor, tokenizer=tokenizer, framework=TensorType.PYTORCH
+            )
             inputs_match, matched_inputs = ensure_model_and_config_inputs_match(model, model_inputs.keys())
             onnx_outputs = list(config.outputs.keys())
 
@@ -150,26 +157,29 @@ def export_pytorch(
 
 
 def export_tensorflow(
-    tokenizer: PreTrainedTokenizer,
+    preprocessor: Union[PreTrainedTokenizer, FeatureExtractionMixin],
     model: TFPreTrainedModel,
     config: OnnxConfig,
     opset: int,
     output: Path,
+    tokenizer: PreTrainedTokenizer = None,
 ) -> Tuple[List[str], List[str]]:
     """
     Export a TensorFlow model to an ONNX Intermediate Representation (IR)
 
     Args:
-        tokenizer ([`PreTrainedTokenizer`]):
-            The tokenizer used for encoding the data.
-        model ([`TFPreTrainedModel`]):
+        preprocessor: ([`PreTrainedTokenizer`] or [`FeatureExtractionMixin`]):
+            The preprocessor used for encoding the data.
+        model ([`PreTrainedModel`] or [`TFPreTrainedModel`]):
             The model to export.
         config ([`~onnx.config.OnnxConfig`]):
             The ONNX configuration associated with the exported model.
         opset (`int`):
             The version of the ONNX operator set to use.
         output (`Path`):
             Directory to store the exported ONNX model.
+        tokenizer ([`PreTrainedTokenizer`]):
+            The tokenizer used for encoding the data.
 
     Returns:
         `Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from
@@ -190,7 +200,7 @@ def export_tensorflow(
             setattr(model.config, override_config_key, override_config_value)
 
     # Ensure inputs match
-    model_inputs = config.generate_dummy_inputs(tokenizer, framework=TensorType.TENSORFLOW)
+    model_inputs = config.generate_dummy_inputs(preprocessor, tokenizer=tokenizer, framework=TensorType.TENSORFLOW)
     inputs_match, matched_inputs = ensure_model_and_config_inputs_match(model, model_inputs.keys())
     onnx_outputs = list(config.outputs.keys())
 
@@ -203,18 +213,19 @@ def export_tensorflow(
 
 
 def export(
-    tokenizer: PreTrainedTokenizer,
+    preprocessor: Union[PreTrainedTokenizer, FeatureExtractionMixin],
     model: Union[PreTrainedModel, TFPreTrainedModel],
     config: OnnxConfig,
     opset: int,
     output: Path,
+    tokenizer: PreTrainedTokenizer = None,
 ) -> Tuple[List[str], List[str]]:
     """
     Export a Pytorch or TensorFlow model to an ONNX Intermediate Representation (IR)
 
     Args:
-        tokenizer ([`PreTrainedTokenizer`]):
-            The tokenizer used for encoding the data.
+        preprocessor: ([`PreTrainedTokenizer`] or [`FeatureExtractionMixin`]):
+            The preprocessor used for encoding the data.
         model ([`PreTrainedModel`] or [`TFPreTrainedModel`]):
             The model to export.
         config ([`~onnx.config.OnnxConfig`]):
@@ -223,6 +234,8 @@ def export(
             The version of the ONNX operator set to use.
         output (`Path`):
             Directory to store the exported ONNX model.
+        tokenizer ([`PreTrainedTokenizer`]):
+            The tokenizer used for encoding the data.
 
     Returns:
         `Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from
@@ -241,18 +254,19 @@ def export(
             raise AssertionError(f"Unsupported PyTorch version, minimum required is 1.8.0, got: {torch_version}")
 
     if is_torch_available() and issubclass(type(model), PreTrainedModel):
-        return export_pytorch(tokenizer, model, config, opset, output)
+        return export_pytorch(preprocessor, model, config, opset, output, tokenizer=tokenizer)
     elif is_tf_available() and issubclass(type(model), TFPreTrainedModel):
-        return export_tensorflow(tokenizer, model, config, opset, output)
+        return export_tensorflow(preprocessor, model, config, opset, output, tokenizer=tokenizer)
 
 
 def validate_model_outputs(
     config: OnnxConfig,
-    tokenizer: PreTrainedTokenizer,
+    preprocessor: Union[PreTrainedTokenizer, FeatureExtractionMixin],
     reference_model: Union[PreTrainedModel, TFPreTrainedModel],
     onnx_model: Path,
     onnx_named_outputs: List[str],
     atol: float,
+    tokenizer: PreTrainedTokenizer = None,
 ):
     from onnxruntime import InferenceSession, SessionOptions
 
@@ -261,9 +275,15 @@ def validate_model_outputs(
     # TODO: generate inputs with a different batch_size and seq_len that was used for conversion to properly test
     # dynamic input shapes.
     if issubclass(type(reference_model), PreTrainedModel):
-        reference_model_inputs = config.generate_dummy_inputs(tokenizer, framework=TensorType.PYTORCH)
+        reference_model_inputs = config.generate_dummy_inputs(
+            preprocessor,
+            tokenizer=tokenizer,
+            framework=TensorType.PYTORCH,
+        )
     else:
-        reference_model_inputs = config.generate_dummy_inputs(tokenizer, framework=TensorType.TENSORFLOW)
+        reference_model_inputs = config.generate_dummy_inputs(
+            preprocessor, tokenizer=tokenizer, framework=TensorType.TENSORFLOW
+        )
 
     # Create ONNX Runtime session
     options = SessionOptions()