Add AutoModel.from_pretrained support (unet3plus+efficientnet)

config.json

@@ -6,7 +6,7 @@
   "out_channels": 1,
   "inter_ch": 64,
   "auto_map": {
-    "AutoConfig": "configuration_unet3plus.UNet3PlusConfig",
+    "AutoConfig": "modeling_unet3plus.UNet3PlusConfig",
     "AutoModel": "modeling_unet3plus.UNet3PlusForSegmentation"
   },
   "training": {

modeling_unet3plus.py

@@ -7,12 +7,31 @@ if _here not in sys.path:
 
 import torch
 import torch.nn.functional as F
-from transformers import PreTrainedModel
+from transformers import PreTrainedModel, PretrainedConfig
 
-from configuration_unet3plus import UNet3PlusConfig
 from models import create_model
 
 
+class UNet3PlusConfig(PretrainedConfig):
+    model_type = "unet3plus"
+
+    def __init__(
+        self,
+        backbone: str = "efficientnet",
+        in_channels: int = 3,
+        out_channels: int = 1,
+        inter_ch: int = 64,
+        img_size: int = 256,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.backbone = backbone
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.inter_ch = inter_ch
+        self.img_size = img_size
+
+
 class UNet3PlusForSegmentation(PreTrainedModel):
     """UNet 3+ segmentation model with EfficientNet-B0 backbone.