Add ONNX export for ViT #15658

lewtun · 2022-02-15T12:42:25Z

What does this PR do?

This PR enables the export of Vision Transformers (ViT) to ONNX with the following features:

default
image-classification

To enable this new modality, I had to significantly refactor the internals of the ONNX exporter because we need a way to pass the feature extractor instead of the tokenizer.

Thanks to a tip from @LysandreJik I replaced the positional tokenizer argument in various functions with a new preprocessor argument that can be a tokenizer or feature extractor (and possibly a processor in future). This should guarantee backwards compatibility for users who chose to use the Python API instead of the transformers.onnx CLI.

Usage

import requests
import numpy as np
from PIL import Image
from onnxruntime import InferenceSession
from transformers import AutoConfig, AutoFeatureExtractor, AutoModelForImageClassification

# Export ViT checkpoint with image classification head
model_ckpt = "google/vit-base-patch16-224"
!python -m transformers.onnx --model={model_ckpt} --feature=image-classification onnx/

# Download an image of two cute cats - naturally ;-)
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)

# Instantiate config and feature extractor
config = AutoConfig.from_pretrained(model_ckpt)
feature_extractor = AutoFeatureExtractor.from_pretrained(model_ckpt)
inputs = feature_extractor(image, return_tensors="np")

# Create ONNX Runtime session
session = InferenceSession("onnx/model.onnx", providers=["CPUExecutionProvider"])
outputs = session.run(["logits"], dict(inputs))
predicted_class_idx = np.argmax(outputs[0])
# Returns Predicted class: Egyptian cat
print("Predicted class:", config.id2label[predicted_class_idx])

Here's two Colab notebooks comparing the inference gains with ORT vs vanilla PyTorch (~20-30% faster on CPU, ~5% faster on GPU):

Todo

Add deprecation warning if user passes tokenizer as keyword argument
Run an inference test to see if we get any speed-up over vanilla PyTorch (maybe)

HuggingFaceDocBuilder · 2022-02-15T12:42:44Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

src/transformers/onnx/config.py

lewtun · 2022-02-15T14:11:29Z

tests/test_onnx_v2.py

@@ -270,6 +276,7 @@ def _onnx_export(self, test_name, name, model_name, feature, onnx_config_class_c
    @parameterized.expand(_get_models_to_test(PYTORCH_EXPORT_MODELS))
    @slow
    @require_torch
+    @require_vision


I added the vision requirement here to test the ViT checkpoint. Please let me know if this isn't a "good practice" because it mixes multiple modalities together

I don't think the vision modality is installed for ONNX tests, so you'd have to double check this actually ends up being tested.

src/transformers/onnx/config.py

src/transformers/models/vit/configuration_vit.py

src/transformers/onnx/config.py

src/transformers/onnx/convert.py

michaelbenayoun · 2022-02-15T15:58:58Z

tests/test_onnx_v2.py

        model = model_class.from_config(config)
        onnx_config = onnx_config_class_constructor(model.config)

+        # Check the modality of the inputs and instantiate the appropriate preprocessor


If this becomes a piece of code we use often, maybe we can refactor this into a function?

lewtun · 2022-02-15T15:59:58Z

src/transformers/onnx/config.py

+            data = np.random.rand(image_height, image_width, num_channels) * 255
+            images.append(Image.fromarray(data.astype("uint8")).convert("RGB"))
+        return images
+
    def generate_dummy_inputs(


This base method now has a mix of arguments for text and image modalities. I'm not 100% sure if we should split the modalities apart ...

You split it now right? Just checking to make sure.

sgugger

Thanks for adding this!

Regarding the the tokenizer optional kwarg, it's very good to keep it like this, but there should be a deprecation warning when it's actually used, and it shouldn't be documented.

src/transformers/onnx/__main__.py

src/transformers/onnx/config.py

src/transformers/onnx/convert.py

sgugger · 2022-02-15T18:13:33Z

tests/test_onnx_v2.py

@@ -270,6 +276,7 @@ def _onnx_export(self, test_name, name, model_name, feature, onnx_config_class_c
    @parameterized.expand(_get_models_to_test(PYTORCH_EXPORT_MODELS))
    @slow
    @require_torch
+    @require_vision


I don't think the vision modality is installed for ONNX tests, so you'd have to double check this actually ends up being tested.

lewtun · 2022-02-17T19:53:07Z

While testing this branch on Colab, I discovered a weird bug when trying to run inference in ONNX Runtime with torch v1.10.2:

RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running Reshape node. Name:'Reshape_42' Status Message: /Users/runner/work/1/s/onnxruntime/core/providers/cpu/tensor/reshape_helper.h:42 onnxruntime::ReshapeHelper::ReshapeHelper(const onnxruntime::TensorShape &, std::vector<int64_t> &, bool) gsl::narrow_cast<int64_t>(input_shape.Size()) == size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{1,197,768}, requested shape:{2,197,12,64}

Curiously, there is no problem running inference with torch v1.9, so something seems to have changed in the torch ONNX exporter in the latest version. I'm currently investigating what the source of the problem is ...

Co-authored-by: Sylvain Gugger <[email protected]>

…onnx-export

lewtun · 2022-02-24T14:30:07Z

src/transformers/onnx/config.py

 from .utils import ParameterFormat, compute_effective_axis_dimension, compute_serialized_parameters_size


+if TYPE_CHECKING:


Since I was already sorting out the relative imports, I also went ahead and fixed the import that are just used for type checking

❤️❤️ ❤️ ❤️

lewtun · 2022-02-24T14:30:57Z

src/transformers/onnx/config.py

-    DEFAULT_FIXED_SEQUENCE = 8
-
-    _TASKS_TO_COMMON_OUTPUTS = {
+    default_fixed_batch = 2


These class variables are now snake_case to prevent confusion / disaster with global constants

sgugger

Thank you so much for making this file more resilient and less prone to cyclical import errors :-)

sgugger · 2022-02-24T14:50:41Z

src/transformers/onnx/convert.py

+if is_torch_available():
+    from ..modeling_utils import PreTrainedModel
+
+if is_tf_available():
+    from ..modeling_tf_utils import TFPreTrainedModel


Thank your for this 😍 !

sgugger · 2022-02-24T14:52:08Z

src/transformers/onnx/config.py


+from ..feature_extraction_utils import FeatureExtractionMixin
+from ..file_utils import TensorType, is_torch_available, is_vision_available
+from ..tokenization_utils_base import PreTrainedTokenizerBase


Last step since this file is imported at very low level, it would be great to import those (PreTrainedTokenizerBase and FeatureExtractionMixin) in TYPE_CHECKING (for type checks) and then only when we do the instance check dynamically

Sounds good!

NielsRogge · 2022-02-28T12:46:44Z

.gitignore

Not sure if we want to add this in the general gitignore of Transformers?

Oop! Will fix that!

michaelbenayoun

Awesome work @lewtun !

michaelbenayoun · 2022-02-28T14:25:56Z

src/transformers/onnx/config.py

+            data = np.random.rand(image_height, image_width, num_channels) * 255
+            images.append(Image.fromarray(data.astype("uint8")).convert("RGB"))
+        return images
+
    def generate_dummy_inputs(


You split it now right? Just checking to make sure.

michaelbenayoun · 2022-02-28T14:27:44Z

src/transformers/onnx/config.py

@@ -326,7 +401,7 @@ def num_attention_heads(self) -> int:

    def generate_dummy_inputs(
        self,
-        tokenizer: PreTrainedTokenizer,
+        tokenizer: "PreTrainedTokenizerBase",


Are you asking about the change to PreTrainedTokenizerBase or use of strings for the typing? Here's the reasons in both cases:

I chose PreTrainedTokenizerBase because it covers both slow and fast tokenizers. The alternative would have been something like Union[PreTrainedTokenizer, PreTrainedTokenizerFast], but that felt clunky

I used strings for the typing following @sgugger's suggestion to use the TYPE_CHECKING constant to fix the circular imports

I was asking about the change of class, and it makes sense to me now, thanks for the explanation!

LysandreJik

Looks good! Thanks @lewtun for iterating and @sgugger for the great reviews!

LysandreJik · 2022-03-09T12:03:38Z

src/transformers/onnx/config.py

@@ -178,6 +196,21 @@ def atol_for_validation(self) -> float:
        """
        return 1e-5

+    @property
+    def is_torch_support_available(self) -> bool:


For torch.fx we have a requirement on a specific torch version. If you have validated that it doesn't work with a specific torch version, I would see no problem in printing a warning mentioning exactly that. If it's going to fail, then raising an error is also fine.

davanstrien · 2022-03-10T07:25:22Z

Super happy to see this merged! 🤗

lewtun added 3 commits February 11, 2022 22:16

Add ONNX support for ViT

8fce819

Refactor to use generic preprocessor

5e15830

Refactor

be90f25

lewtun added 5 commits February 15, 2022 13:47

Fix ONNX conversion for models with fast tokenizers

81287ec

Fix copies

8103587

Add vision to tests

bcbdbd9

Remove fixed ViT outputs

12a5306

Extend ONNX slow tests to ViT

ba0a7b0

lewtun commented Feb 15, 2022

View reviewed changes

src/transformers/onnx/config.py Outdated Show resolved Hide resolved

lewtun changed the title ~~Add ONNX export for vision models~~ Add ONNX export for ViT Feb 15, 2022

lewtun requested review from NielsRogge, michaelbenayoun, sgugger and LysandreJik February 15, 2022 14:06

lewtun commented Feb 15, 2022

View reviewed changes

src/transformers/onnx/config.py Outdated Show resolved Hide resolved

lewtun commented Feb 15, 2022

View reviewed changes

src/transformers/onnx/config.py Outdated Show resolved Hide resolved

Add dummy image generator

0ebbcae

michaelbenayoun reviewed Feb 15, 2022

View reviewed changes

lewtun commented Feb 15, 2022

View reviewed changes

sgugger reviewed Feb 15, 2022

View reviewed changes

lewtun added 5 commits February 15, 2022 21:27

Use model_type to determine modality

d179861

Add deprecation warnings for tokenizer argument

b1b4f61

Add warning when overwriting the preprocessor

b0491e8

Add optional args to docstrings

7f03d43

Add TODO

26dcdde

lewtun added 2 commits February 23, 2022 17:14

Add minimum PyTorch version to OnnxConfig

99dd9a7

Merge branch 'master' into vision-onnx-export

41fa7e0

lewtun and others added 9 commits February 24, 2022 12:21

Replace absolute imports with relative ones

ade513f

Apply Sylvain's suggestions from code review

a7baf9a

Co-authored-by: Sylvain Gugger <[email protected]>

Merge remote-tracking branch 'origin/vision-onnx-export' into vision-…

d898037

…onnx-export

Fix imports

88d25cf

Refactor OnnxConfig class variables from CONSTANT_NAME to snake_case

941689b

Fix ViT torch version

31dd4f9

Fix docstring

d1f9397

Fix imports and add logging

49dca94

Use relative imports for real this time and use type checking

aec42f8

lewtun commented Feb 24, 2022

View reviewed changes

Add check for vision feature extractor

750db82

sgugger approved these changes Feb 24, 2022

View reviewed changes

lewtun added 3 commits February 24, 2022 16:36

Refactor imports for type checking

951df50

Skip ONNX test if torch version is incompatible

48129a7

Revert ImportError vs AssertionError

b2e618e

NielsRogge reviewed Feb 28, 2022

View reviewed changes

Revert gitignore

1514807

michaelbenayoun approved these changes Feb 28, 2022

View reviewed changes

lewtun mentioned this pull request Mar 7, 2022

GeneratorExp aren't supported by torch.jit.script when I try to export a previously trained model 'google/vit-base-patch16-224-in21k'. #15354

Open

1 task

LysandreJik approved these changes Mar 9, 2022

View reviewed changes

lewtun added 2 commits March 9, 2022 14:31

Replace ImportError with warning

7ac6312

Add reasonable value for default atol

81eedea

lewtun merged commit 50dd314 into master Mar 9, 2022

lewtun deleted the vision-onnx-export branch March 9, 2022 16:37

This was referenced Mar 10, 2022

Fix duplicate arguments passed to dummy inputs in ONNX export #16045

Merged

Enable ONNX export for VisionDecoderEncoderModel #14812

Closed

akuma12 mentioned this pull request Mar 30, 2022

Add ONNX export for BeiT #16498

Merged

mht-sharma mentioned this pull request Oct 12, 2022

Add support ORT whisper huggingface/optimum#420

Merged

3 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add ONNX export for ViT #15658

Add ONNX export for ViT #15658

lewtun commented Feb 15, 2022 •

edited

Loading

HuggingFaceDocBuilder commented Feb 15, 2022

lewtun Feb 15, 2022

sgugger Feb 15, 2022

michaelbenayoun Feb 15, 2022

lewtun Feb 15, 2022

michaelbenayoun Feb 28, 2022

sgugger left a comment

sgugger Feb 15, 2022

lewtun commented Feb 17, 2022

lewtun Feb 24, 2022

sgugger Feb 24, 2022

lewtun Feb 24, 2022

sgugger left a comment

sgugger Feb 24, 2022

sgugger Feb 24, 2022

lewtun Feb 24, 2022

NielsRogge Feb 28, 2022

lewtun Feb 28, 2022

lewtun Feb 28, 2022

michaelbenayoun left a comment

michaelbenayoun Feb 28, 2022

michaelbenayoun Feb 28, 2022

lewtun Feb 28, 2022

michaelbenayoun Feb 28, 2022

LysandreJik left a comment

LysandreJik Mar 9, 2022

davanstrien commented Mar 10, 2022

		from .utils import ParameterFormat, compute_effective_axis_dimension, compute_serialized_parameters_size


		if TYPE_CHECKING:

Add ONNX export for ViT #15658

Add ONNX export for ViT #15658

Conversation

lewtun commented Feb 15, 2022 • edited Loading

What does this PR do?

Usage

Todo

HuggingFaceDocBuilder commented Feb 15, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sgugger left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

lewtun commented Feb 17, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sgugger left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

michaelbenayoun left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

LysandreJik left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

davanstrien commented Mar 10, 2022

lewtun commented Feb 15, 2022 •

edited

Loading