File size: 518 Bytes
c2c83fe 3283c80 c2c83fe 3283c80 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
from transformers import PretrainedConfig
from typing import Dict
class MobileNetV3Config(PretrainedConfig):
model_type = "mobilenetv3"
def __init__(
self,
num_classes: int=6,
id2label: Dict={
0: "cardboard",
1: "glass",
2: "metal",
3: "paper",
4: "plastic",
5: "trash"
},
**kwargs,
):
super().__init__(**kwargs)
self.id2label = id2label
self.num_classes = num_classes |