SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

[MobileNet-v2] Fix ONNX typo (#20860)

Browse Source 
* fix typo `onnx`

* fix test
This commit is contained in:
Younes Belkada
2022-12-22 18:52:54 +01:00
committed by GitHub
parent 4d10ffd506
commit 52dd2b61bf
2 changed files with 8 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/transformers/onnx/features.py
View File

@@ -408,12 +408,12 @@ class FeaturesManager:
"question-answering",
onnx_config_cls="models.mobilebert.MobileBertOnnxConfig",
),
"mobilenet_v1": supported_features_mapping(
"mobilenet-v1": supported_features_mapping(
"default",
"image-classification",
onnx_config_cls="models.mobilenet_v1.MobileNetV1OnnxConfig",
),
"mobilenet_v2": supported_features_mapping(
"mobilenet-v2": supported_features_mapping(
"default",
"image-classification",
onnx_config_cls="models.mobilenet_v2.MobileNetV2OnnxConfig",

7
tests/onnx/test_onnx_v2.py
View File

@@ -272,7 +272,12 @@ def _get_models_to_test(export_models_list):
feature: FeaturesManager.get_config(name, feature) for _ in features for feature in _
}
else:
feature_config_mapping = FeaturesManager.get_supported_features_for_model_type(name)
# pre-process the model names
model_type = name.replace("_", "-")
model_name = getattr(model, "name", "")
feature_config_mapping = FeaturesManager.get_supported_features_for_model_type(
model_type, model_name=model_name
)
for feature, onnx_config_class_constructor in feature_config_mapping.items():
models_to_test.append((f"{name}_{feature}", name, model, feature, onnx_config_class_constructor))