|
|
from typing import Dict
|
|
|
|
|
|
from transformers.configuration_utils import PretrainedConfig
|
|
|
from optimum.exporters.onnx.model_configs import ViTOnnxConfig
|
|
|
|
|
|
MODEL_NAMES = [
|
|
|
'efficientnet_b0',
|
|
|
'efficientnet_b1',
|
|
|
'efficientnet_b2',
|
|
|
'efficientnet_b3',
|
|
|
'efficientnet_b4',
|
|
|
'efficientnet_b5',
|
|
|
'efficientnet_b6',
|
|
|
'efficientnet_b7',
|
|
|
'efficientnet_b8',
|
|
|
'efficientnet_l2'
|
|
|
]
|
|
|
|
|
|
|
|
|
class EfficientNetConfig(PretrainedConfig):
|
|
|
model_type = 'efficientnet'
|
|
|
|
|
|
def __init__(
|
|
|
self,
|
|
|
model_name: str = 'efficientnet_b0',
|
|
|
pretrained: bool = False,
|
|
|
num_classes: int = 1000,
|
|
|
global_pool: str = 'avg',
|
|
|
**kwargs,
|
|
|
):
|
|
|
if model_name not in MODEL_NAMES:
|
|
|
raise ValueError(f'`model_name` must be one of these: {MODEL_NAMES}, but got {model_name}')
|
|
|
|
|
|
self.model_name = model_name
|
|
|
self.pretrained = pretrained
|
|
|
self.num_classes = num_classes
|
|
|
self.global_pool = global_pool
|
|
|
super().__init__(**kwargs)
|
|
|
|
|
|
|
|
|
class EfficientNetOnnxConfig(ViTOnnxConfig):
|
|
|
@property
|
|
|
def outputs(self) -> Dict[str, Dict[int, str]]:
|
|
|
common_outputs = super().outputs
|
|
|
|
|
|
if self.task == "image-classification":
|
|
|
common_outputs["logits"] = {0: "batch_size", 1: "num_classes"}
|
|
|
|
|
|
return common_outputs
|
|
|
|
|
|
|
|
|
__all__ = [
|
|
|
'EfficientNetConfig',
|
|
|
'EfficientNetOnnxConfig'
|
|
|
] |
