From fe785730dcbf3390aa07f667e8d3c4b02d6638e0 Mon Sep 17 00:00:00 2001 From: Niklas Hansson Date: Tue, 9 Aug 2022 09:35:05 +0200 Subject: [PATCH] fix: data2vec-vision Onnx ready-made configuration. (#18427) * feat: add the data2vec conf that are missing https://huggingface.co/docs/transformers/serialization * fix: wrong config --- src/transformers/onnx/features.py | 6 ++++++ tests/onnx/test_onnx_v2.py | 1 + 2 files changed, 7 insertions(+) diff --git a/src/transformers/onnx/features.py b/src/transformers/onnx/features.py index e7c24a8ad9..3eea94c8c1 100644 --- a/src/transformers/onnx/features.py +++ b/src/transformers/onnx/features.py @@ -229,6 +229,12 @@ class FeaturesManager: "question-answering", onnx_config_cls="models.data2vec.Data2VecTextOnnxConfig", ), + "data2vec-vision": supported_features_mapping( + "default", + "image-classification", + "image-segmentation", + onnx_config_cls="models.data2vec.Data2VecVisionOnnxConfig", + ), "deberta": supported_features_mapping( "default", "masked-lm", diff --git a/tests/onnx/test_onnx_v2.py b/tests/onnx/test_onnx_v2.py index cfc58dd335..c22406841a 100644 --- a/tests/onnx/test_onnx_v2.py +++ b/tests/onnx/test_onnx_v2.py @@ -207,6 +207,7 @@ PYTORCH_EXPORT_MODELS = { ("deit", "facebook/deit-small-patch16-224"), ("beit", "microsoft/beit-base-patch16-224"), ("data2vec-text", "facebook/data2vec-text-base"), + ("data2vec-vision", "facebook/data2vec-vision-base"), ("perceiver", "deepmind/language-perceiver", ("masked-lm", "sequence-classification")), ("perceiver", "deepmind/vision-perceiver-conv", ("image-classification",)), ("yolos", "hustvl/yolos-tiny"),