Upload model

config.json

@@ -0,0 +1,12 @@
+{
+  "architectures": [
+    "BBSNetModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_bbsnet.BBSNetConfig",
+    "AutoModel": "modeling_bbsnet.BBSNetModel"
+  },
+  "model_type": "bbsnet",
+  "torch_dtype": "float32",
+  "transformers_version": "4.26.1"
+}

configuration_bbsnet.py

@@ -0,0 +1,47 @@
+from typing import List
+
+from transformers import PretrainedConfig
+
+"""
+The configuration of a model is an object that 
+will contain all the necessary information to build the model.
+
+The three important things to remember when writing you own configuration are the following:
+
+- you have to inherit from PretrainedConfig,
+- the __init__ of your PretrainedConfig must accept any kwargs,
+- those kwargs need to be passed to the superclass __init__.
+"""
+
+
+class BBSNetConfig(PretrainedConfig):
+
+    """
+    Defining a model_type for your configuration is not mandatory,
+    unless you want to register your model with the auto classes."""
+
+    model_type = "bbsnet"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+
+if __name__ == "__main__":
+    """
+    With this done, you can easily create and save your configuration like
+    you would do with any other model config of the library.
+    Here is how we can create a resnet50d config and save it:
+    """
+    bbsnet_config = BBSNetConfig()
+    bbsnet_config.save_pretrained("custom-bbsnet")
+
+    """
+    This will save a file named config.json inside the folder custom-resnet. 
+    You can then reload your config with the from_pretrained method:
+    """
+    bbsnet_config = BBSNetConfig.from_pretrained("custom-bbsnet")
+
+    """
+    You can also use any other method of the PretrainedConfig class, 
+    like push_to_hub() to directly upload your config to the Hub.
+    """

modeling_bbsnet.py

@@ -0,0 +1,48 @@
+from typing import Dict, Optional
+
+from torch import Tensor, nn
+from transformers import PreTrainedModel
+
+from models.BBSNet_model import BBSNet
+
+from .configuration_bbsnet import BBSNetConfig
+
+
+class BBSNetModel(PreTrainedModel):
+    """
+    The line that sets the config_class is not mandatory,
+    unless you want to register your model with the auto classes
+    """
+
+    config_class = BBSNetConfig
+
+    def __init__(self, config: BBSNetConfig):
+        super().__init__(config)
+        self.model = BBSNet()
+        self.loss = nn.BCEWithLogitsLoss()
+
+    """
+    You can have your model return anything you want, 
+    but returning a dictionary with the loss included when labels are passed, 
+    will make your model directly usable inside the Trainer class. 
+    Using another output format is fine as long as you are planning on 
+    using your own training loop or another library for training.
+    """
+
+    def forward(
+        self, rgbs: Tensor, depths: Tensor, gts: Optional[Tensor] = None
+    ) -> Dict[str, Tensor]:
+        _, logits = self.model(rgbs, depths)
+        if gts is not None:
+            loss = self.loss(logits, gts)
+            return {"loss": loss, "logits": logits}
+        return {"logits": logits}
+
+
+if __name__ == "__main__":
+    resnet50d_config = ResnetConfig.from_pretrained("custom-resnet")
+    resnet50d = ResnetModelForImageClassification(resnet50d_config)
+
+    # Load pretrained weights from timm
+    pretrained_model: nn.Module = timm.create_model("resnet50d", pretrained=True)
+    resnet50d.model.load_state_dict(pretrained_model.state_dict())

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:483479ab2cc3cd42ff1be16e1ba76ac09afcb0c010487221b81c96aa22f75106
+size 199976498