Add dpt-hybrid support

src/transformers/modeling_outputs.py

@@ -65,6 +65,59 @@ class BaseModelOutputWithNoAttention(ModelOutput):
     hidden_states: Optional[Tuple[torch.FloatTensor]] = None
 
 
+@dataclass
+class BaseModelOutputWithIntermediateActivations(ModelOutput):
+    """
+    Base class for model's outputs that also contains intermediate activations that can be used at later stages. Useful
+    in the context of Vision models.:
+
+    Args:
+        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        intermediate_activations (`tuple(torch.FloatTensor)`, *optional*):
+            Intermediate activations that can be used to compute hidden states of the model at various layers.
+    """
+
+    last_hidden_states: torch.FloatTensor = None
+    intermediate_activations: Optional[Tuple[torch.FloatTensor]] = None
+
+
+@dataclass
+class BaseModelOutputWithPoolingAndIntermediateActivations(ModelOutput):
+    """
+    Base class for model's outputs that also contains a pooling of the last hidden states as well as intermediate
+    activations that can be used by the model at later stages.
+
+    Args:
+        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`):
+            Last layer hidden-state of the first token of the sequence (classification token) after further processing
+            through the layers used for the auxiliary pretraining task. E.g. for BERT-family of models, this returns
+            the classification token after processing through a linear layer and a tanh activation function. The linear
+            layer weights are trained from the next sentence prediction (classification) objective during pretraining.
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+        intermediate_activations (`tuple(torch.FloatTensor)`, *optional*):
+            Intermediate activations that can be used to compute hidden states of the model at various layers.
+    """
+
+    last_hidden_state: torch.FloatTensor = None
+    pooler_output: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+    intermediate_activations: Optional[Tuple[torch.FloatTensor]] = None
+
+
 @dataclass
 class BaseModelOutputWithPooling(ModelOutput):
     """

src/transformers/models/dpt/configuration_dpt.py

@@ -14,8 +14,11 @@
 # limitations under the License.
 """ DPT model configuration"""
 
+import copy
+
 from ...configuration_utils import PretrainedConfig
 from ...utils import logging
+from ..bit import BitConfig
 
 
 logger = logging.get_logger(__name__)
@@ -76,6 +79,8 @@ class DPTConfig(PretrainedConfig):
             - "project" passes information to the other tokens by concatenating the readout to all other tokens before
               projecting the
             representation to the original feature dimension D using a linear layer followed by a GELU non-linearity.
+        embedding_type (`str`, *optional*, defaults to `"patch_embedding"`):
+            The type of embedding to use. Can be one of [`"patch_embedding"`, `"hybrid"`].
         reassemble_factors (`List[int]`, *optional*, defaults to `[4, 2, 1, 0.5]`):
             The up/downsampling factors of the reassemble layers.
         neck_hidden_sizes (`List[str]`, *optional*, defaults to [96, 192, 384, 768]):
@@ -94,6 +99,12 @@ class DPTConfig(PretrainedConfig):
             The index that is ignored by the loss function of the semantic segmentation model.
         semantic_classifier_dropout (`float`, *optional*, defaults to 0.1):
             The dropout ratio for the semantic classification head.
+        backbone_featmap_shape (`List[int]`, *optional*, defaults to `[1, 1024, 24, 24]`):
+            Used only for the `hybrid` embedding type. The shape of the feature maps of the backbone.
+        neck_ignore_stages (`List[int]`, *optional*, defaults to `[0, 1]`):
+            Used only for the `hybrid` embedding type. The stages of the readout layers to ignore.
+        backbone_config (`Dict[str, Any]`, *optional*, defaults to `None`):
+            Used only for the `hybrid` embedding type. The configuration of the backbone in a dictionary.
 
     Example:
 
@@ -125,6 +136,7 @@ def __init__(
         image_size=384,
         patch_size=16,
         num_channels=3,
+        embedding_type="patch_embedding",
         qkv_bias=True,
         backbone_out_indices=[2, 5, 8, 11],
         readout_type="project",
@@ -137,11 +149,43 @@ def __init__(
         auxiliary_loss_weight=0.4,
         semantic_loss_ignore_index=255,
         semantic_classifier_dropout=0.1,
+        backbone_featmap_shape=[1, 1024, 24, 24],
+        neck_ignore_stages=[0, 1],
+        backbone_config=None,
         **kwargs
     ):
         super().__init__(**kwargs)
 
         self.hidden_size = hidden_size
+
+        if embedding_type not in ["patch_embedding", "hybrid"]:
+            raise ValueError("Embedding type must be one of ['patch_embedding', 'hybrid']")
+        if embedding_type == "hybrid":
+            logger.info("Initializing the config with a `BiT` backbone.")
+            if backbone_config is None:
+                backbone_config = {
+                    "global_padding": "same",
+                    "layer_type": "bottleneck",
+                    "depths": [3, 4, 9],
+                    "out_features": ["stage1", "stage2", "stage3"],
+                    "embedding_dynamic_padding": True,
+                }
+            elif not isinstance(backbone_config, dict):
+                raise ValueError("backbone_config must be a dictionary.")
+            self.backbone_config = BitConfig(**backbone_config)
+            self.is_hybrid = True
+            self.backbone_featmap_shape = backbone_featmap_shape
+            self.neck_ignore_stages = neck_ignore_stages
+
+            if readout_type != "project":
+                raise ValueError("Readout type must be 'project' when using `DPT-hybrid` mode.")
+        else:
+            self.backbone_config = None
+            self.is_hybrid = False
+            self.backbone_featmap_shape = None
+            self.neck_ignore_stages = []
+        self.embedding_type = embedding_type
+
         self.num_hidden_layers = num_hidden_layers
         self.num_attention_heads = num_attention_heads
         self.intermediate_size = intermediate_size
@@ -168,3 +212,16 @@ def __init__(
         self.auxiliary_loss_weight = auxiliary_loss_weight
         self.semantic_loss_ignore_index = semantic_loss_ignore_index
         self.semantic_classifier_dropout = semantic_classifier_dropout
+
+    def to_dict(self):
+        """
+        Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`]. Returns:
+            `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
+        """
+        output = copy.deepcopy(self.__dict__)
+
+        if output["backbone_config"] is not None:
+            output["backbone_config"] = self.backbone_config.to_dict()
+
+        output["model_type"] = self.__class__.model_type
+        return output