Upload model

config.json

@@ -9,7 +9,6 @@
   "kernel_size": 3,
   "model_type": "basnet",
   "n_channels": 3,
-  "resnet_model": "microsoft/resnet-34",
   "torch_dtype": "float32",
   "transformers_version": "4.42.4"
 }

configuration_basnet.py

@@ -6,13 +6,10 @@ class BASNetConfig(PretrainedConfig):
 
     def __init__(
         self,
-        resnet_model: str = "microsoft/resnet-34",
         n_channels: int = 3,
         kernel_size: int = 3,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
-        self.resnet_model = resnet_model
         self.n_channels = n_channels
-
-        self.kernel_size = 3
+        self.kernel_size = kernel_size

modeling_basnet.py

@@ -1,16 +1,30 @@
 import logging
-from typing import Optional, Tuple
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
 import torchvision
 from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import ModelOutput
 
 from .configuration_basnet import BASNetConfig
 
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class BASNetModelOutput(ModelOutput):
+    dout: torch.Tensor
+    d1: Optional[torch.Tensor] = None
+    d2: Optional[torch.Tensor] = None
+    d3: Optional[torch.Tensor] = None
+    d4: Optional[torch.Tensor] = None
+    d5: Optional[torch.Tensor] = None
+    d6: Optional[torch.Tensor] = None
+    db: Optional[torch.Tensor] = None
+
+
 class RefUnet(nn.Module):
     def __init__(self, in_ch: int, inc_ch: int) -> None:
         super().__init__()
@@ -352,17 +366,8 @@ class BASNetModel(PreTrainedModel):
         self.post_init()
 
     def forward(
-        self, pixel_values: torch.Tensor
-    ) -> Tuple[
-        torch.Tensor,
-        torch.Tensor,
-        torch.Tensor,
-        torch.Tensor,
-        torch.Tensor,
-        torch.Tensor,
-        torch.Tensor,
-        torch.Tensor,
-    ]:
+        self, pixel_values: torch.Tensor, return_dict: Optional[bool] = None
+    ) -> Union[Tuple, BASNetModelOutput]:
         hx = pixel_values
 
         ## -------------Encoder-------------
@@ -452,15 +457,36 @@ class BASNetModel(PreTrainedModel):
         ## -------------Refine Module-------------
         dout = self.refunet(d1)  # 256
 
-        return (
-            torch.sigmoid(dout),
-            torch.sigmoid(d1),
-            torch.sigmoid(d2),
-            torch.sigmoid(d3),
-            torch.sigmoid(d4),
-            torch.sigmoid(d5),
-            torch.sigmoid(d6),
-            torch.sigmoid(db),
+        dout_act = torch.sigmoid(dout)
+        d1_act = torch.sigmoid(d1)
+        d2_act = torch.sigmoid(d2)
+        d3_act = torch.sigmoid(d3)
+        d4_act = torch.sigmoid(d4)
+        d5_act = torch.sigmoid(d5)
+        d6_act = torch.sigmoid(d6)
+        db_act = torch.sigmoid(db)
+
+        if not return_dict:
+            return (
+                dout_act,
+                d1_act,
+                d2_act,
+                d3_act,
+                d4_act,
+                d5_act,
+                d6_act,
+                db_act,
+            )
+
+        return BASNetModelOutput(
+            dout=dout_act,
+            d1=d1_act,
+            d2=d2_act,
+            d3=d3_act,
+            d4=d4_act,
+            d5=d5_act,
+            d6=d6_act,
+            db=db_act,
         )