Skip to content

Commit 66e30ad

Browse files
echarlaixecharlaix
authored
Fix tflite export (#1767)
* Set image-feature-extraction as synonym of feature-extraction * format
1 parent 1146eae commit 66e30ad

File tree

2 files changed

+4
-25
lines changed

2 files changed

+4
-25
lines changed

‎optimum/exporters/onnx/base.py

-1
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,6 @@ class OnnxConfig(ExportConfig, ABC):
151151
"feature-extraction": OrderedDict({"last_hidden_state": {0: "batch_size", 1: "sequence_length"}}),
152152
"fill-mask": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}),
153153
"image-classification": OrderedDict({"logits": {0: "batch_size"}}),
154-
"image-feature-extraction": OrderedDict({"last_hidden_state": {0: "batch_size", 1: "sequence_length"}}),
155154
"image-segmentation": OrderedDict({"logits": {0: "batch_size", 1: "num_labels", 2: "height", 3: "width"}}),
156155
"image-to-text": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}),
157156
"image-to-image": OrderedDict(

‎optimum/exporters/tasks.py

+4-24
Original file line numberDiff line numberDiff line change
@@ -168,7 +168,6 @@ class TasksManager:
168168
"feature-extraction": "AutoModel",
169169
"fill-mask": "AutoModelForMaskedLM",
170170
"image-classification": "AutoModelForImageClassification",
171-
"image-feature-extraction": "AutoModel",
172171
"image-segmentation": ("AutoModelForImageSegmentation", "AutoModelForSemanticSegmentation"),
173172
"image-to-image": "AutoModelForImageToImage",
174173
"image-to-text": "AutoModelForVision2Seq",
@@ -257,6 +256,7 @@ class TasksManager:
257256
"translation": "text2text-generation",
258257
"vision2seq-lm": "image-to-text",
259258
"zero-shot-classification": "text-classification",
259+
"image-feature-extraction": "feature-extraction",
260260
}
261261

262262
# Reverse dictionaries str -> str, where several model loaders may map to the same task
@@ -463,13 +463,11 @@ class TasksManager:
463463
),
464464
"convnext": supported_tasks_mapping(
465465
"feature-extraction",
466-
"image-feature-extraction",
467466
"image-classification",
468467
onnx="ConvNextOnnxConfig",
469468
),
470469
"convnextv2": supported_tasks_mapping(
471470
"feature-extraction",
472-
"image-feature-extraction",
473471
"image-classification",
474472
onnx="ConvNextV2OnnxConfig",
475473
),
@@ -486,7 +484,6 @@ class TasksManager:
486484
"data2vec-vision": supported_tasks_mapping(
487485
"feature-extraction",
488486
"image-classification",
489-
"image-feature-extraction",
490487
# ONNX doesn't support `adaptive_avg_pool2d` yet
491488
# "semantic-segmentation",
492489
onnx="Data2VecVisionOnnxConfig",
@@ -520,15 +517,13 @@ class TasksManager:
520517
),
521518
"deit": supported_tasks_mapping(
522519
"feature-extraction",
523-
"image-feature-extraction",
524520
"image-classification",
525521
"masked-im",
526522
onnx="DeiTOnnxConfig",
527523
),
528524
"detr": supported_tasks_mapping(
529525
"feature-extraction",
530526
"object-detection",
531-
"image-feature-extraction",
532527
"image-segmentation",
533528
onnx="DetrOnnxConfig",
534529
),
@@ -555,7 +550,6 @@ class TasksManager:
555550
),
556551
"dpt": supported_tasks_mapping(
557552
"feature-extraction",
558-
"image-feature-extraction",
559553
"depth-estimation",
560554
"image-segmentation",
561555
"semantic-segmentation",
@@ -614,7 +608,6 @@ class TasksManager:
614608
),
615609
"glpn": supported_tasks_mapping(
616610
"feature-extraction",
617-
"image-feature-extraction",
618611
"depth-estimation",
619612
onnx="GlpnOnnxConfig",
620613
),
@@ -682,7 +675,6 @@ class TasksManager:
682675
),
683676
"imagegpt": supported_tasks_mapping(
684677
"feature-extraction",
685-
"image-feature-extraction",
686678
"image-classification",
687679
onnx="ImageGPTOnnxConfig",
688680
),
@@ -714,9 +706,7 @@ class TasksManager:
714706
"token-classification",
715707
onnx="LiltOnnxConfig",
716708
),
717-
"levit": supported_tasks_mapping(
718-
"feature-extraction", "image-classification", "image-feature-extraction", onnx="LevitOnnxConfig"
719-
),
709+
"levit": supported_tasks_mapping("feature-extraction", "image-classification", onnx="LevitOnnxConfig"),
720710
"longt5": supported_tasks_mapping(
721711
"feature-extraction",
722712
"feature-extraction-with-past",
@@ -780,19 +770,17 @@ class TasksManager:
780770
"mobilevit": supported_tasks_mapping(
781771
"feature-extraction",
782772
"image-classification",
783-
"image-feature-extraction",
784773
"image-segmentation",
785774
onnx="MobileViTOnnxConfig",
786775
),
787776
"mobilenet-v1": supported_tasks_mapping(
788777
"feature-extraction",
789-
"image-feature-extraction",
790778
"image-classification",
791779
onnx="MobileNetV1OnnxConfig",
792780
),
793781
"mobilenet-v2": supported_tasks_mapping(
782+
"feature-extraction",
794783
"image-classification",
795-
"image-feature-extraction",
796784
onnx="MobileNetV2OnnxConfig",
797785
),
798786
"mpnet": supported_tasks_mapping(
@@ -901,20 +889,17 @@ class TasksManager:
901889
),
902890
"poolformer": supported_tasks_mapping(
903891
"feature-extraction",
904-
"image-feature-extraction",
905892
"image-classification",
906893
onnx="PoolFormerOnnxConfig",
907894
),
908895
"regnet": supported_tasks_mapping(
909896
"feature-extraction",
910-
"image-feature-extraction",
911897
"image-classification",
912898
onnx="RegNetOnnxConfig",
913899
),
914900
"resnet": supported_tasks_mapping(
915901
"feature-extraction",
916902
"image-classification",
917-
"image-feature-extraction",
918903
onnx="ResNetOnnxConfig",
919904
tflite="ResNetTFLiteConfig",
920905
),
@@ -950,7 +935,6 @@ class TasksManager:
950935
"segformer": supported_tasks_mapping(
951936
"feature-extraction",
952937
"image-classification",
953-
"image-feature-extraction",
954938
"image-segmentation",
955939
"semantic-segmentation",
956940
onnx="SegformerOnnxConfig",
@@ -995,14 +979,12 @@ class TasksManager:
995979
),
996980
"swin": supported_tasks_mapping(
997981
"feature-extraction",
998-
"image-feature-extraction",
999982
"image-classification",
1000983
"masked-im",
1001984
onnx="SwinOnnxConfig",
1002985
),
1003986
"swin2sr": supported_tasks_mapping(
1004987
"feature-extraction",
1005-
"image-feature-extraction",
1006988
"image-to-image",
1007989
onnx="Swin2srOnnxConfig",
1008990
),
@@ -1015,7 +997,6 @@ class TasksManager:
1015997
),
1016998
"table-transformer": supported_tasks_mapping(
1017999
"feature-extraction",
1018-
"image-feature-extraction",
10191000
"object-detection",
10201001
onnx="TableTransformerOnnxConfig",
10211002
),
@@ -1048,7 +1029,7 @@ class TasksManager:
10481029
onnx="VisionEncoderDecoderOnnxConfig",
10491030
),
10501031
"vit": supported_tasks_mapping(
1051-
"feature-extraction", "image-classification", "image-feature-extraction", "masked-im", onnx="ViTOnnxConfig"
1032+
"feature-extraction", "image-classification", "masked-im", onnx="ViTOnnxConfig"
10521033
),
10531034
"wavlm": supported_tasks_mapping(
10541035
"feature-extraction",
@@ -1108,7 +1089,6 @@ class TasksManager:
11081089
),
11091090
"yolos": supported_tasks_mapping(
11101091
"feature-extraction",
1111-
"image-feature-extraction",
11121092
"object-detection",
11131093
onnx="YolosOnnxConfig",
11141094
),

0 commit comments

Comments
 (0)