Add ONNX export for BeiT (#16498)

* Add beit onnx conversion support

* Updated docs

* Added cross reference to ViT ONNX config
This commit is contained in:
Jim Rohrer
2022-04-01 03:52:42 -05:00
committed by GitHub
parent bfeff6cc6a
commit 9de70f213e
5 changed files with 31 additions and 5 deletions

View File

@@ -15,14 +15,13 @@ from transformers.onnx import (
export,
validate_model_outputs,
)
from transformers.onnx.utils import compute_effective_axis_dimension, compute_serialized_parameters_size
from transformers.testing_utils import require_onnx, require_tf, require_torch, require_vision, slow
if is_torch_available() or is_tf_available():
from transformers.onnx.features import FeaturesManager
from transformers.onnx.utils import compute_effective_axis_dimension, compute_serialized_parameters_size
from transformers.testing_utils import require_onnx, require_tf, require_torch, require_vision, slow
@require_onnx
class OnnxUtilsTestCaseV2(TestCase):
@@ -181,6 +180,7 @@ PYTORCH_EXPORT_MODELS = {
("xlm-roberta", "xlm-roberta-base"),
("layoutlm", "microsoft/layoutlm-base-uncased"),
("vit", "google/vit-base-patch16-224"),
("beit", "microsoft/beit-base-patch16-224"),
}
PYTORCH_EXPORT_WITH_PAST_MODELS = {