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,
  )