# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/conditional_detr/modular_conditional_detr.py. # Do NOT edit this file manually as any edits will be overwritten by the generation of # the file from the modular. If any change should be done, please apply the change to the # modular_conditional_detr.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # Copyright 2022 Microsoft Research Asia and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import math from collections.abc import Callable from dataclasses import dataclass import torch from torch import nn from ... import initialization as init from ...activations import ACT2FN from ...backbone_utils import load_backbone from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithCrossAttentions, Seq2SeqModelOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...pytorch_utils import compile_compatible_method_lru_cache from ...utils import ModelOutput, TransformersKwargs, auto_docstring from ...utils.generic import can_return_tuple, merge_with_config_defaults from ...utils.output_capturing import OutputRecorder, capture_outputs from .configuration_conditional_detr import ConditionalDetrConfig @dataclass @auto_docstring( custom_intro=""" Base class for outputs of the CONDITIONAL_DETR decoder. This class adds one attribute to BaseModelOutputWithCrossAttentions, namely an optional stack of intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a layernorm. This is useful when training the model with auxiliary decoding losses. """ ) class ConditionalDetrDecoderOutput(BaseModelOutputWithCrossAttentions): r""" cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=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 of the decoder's cross-attention layer, after the attention softmax, used to compute the weighted average in the cross-attention heads. intermediate_hidden_states (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, num_queries, hidden_size)`, *optional*, returned when `config.auxiliary_loss=True`): Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a layernorm. reference_points (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, num_queries, 2 (anchor points))`): Reference points (reference points of each layer of the decoder). """ intermediate_hidden_states: torch.FloatTensor | None = None reference_points: tuple[torch.FloatTensor] | None = None @dataclass @auto_docstring( custom_intro=""" Base class for outputs of the CONDITIONAL_DETR encoder-decoder model. This class adds one attribute to Seq2SeqModelOutput, namely an optional stack of intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a layernorm. This is useful when training the model with auxiliary decoding losses. """ ) class ConditionalDetrModelOutput(Seq2SeqModelOutput): r""" 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 decoder of the model. intermediate_hidden_states (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, sequence_length, hidden_size)`, *optional*, returned when `config.auxiliary_loss=True`): Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a layernorm. reference_points (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, num_queries, 2 (anchor points))`): Reference points (reference points of each layer of the decoder). """ intermediate_hidden_states: torch.FloatTensor | None = None reference_points: tuple[torch.FloatTensor] | None = None @dataclass @auto_docstring( custom_intro=""" Output type of [`ConditionalDetrForObjectDetection`]. """ ) class ConditionalDetrObjectDetectionOutput(ModelOutput): r""" loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` are provided)): Total loss as a linear combination of a negative log-likehood (cross-entropy) for class prediction and a bounding box loss. The latter is defined as a linear combination of the L1 loss and the generalized scale-invariant IoU loss. loss_dict (`Dict`, *optional*): A dictionary containing the individual losses. Useful for logging. logits (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes + 1)`): Classification logits (including no-object) for all queries. pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`): Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These values are normalized in [0, 1], relative to the size of each individual image in the batch (disregarding possible padding). You can use [`~ConditionalDetrImageProcessor.post_process_object_detection`] to retrieve the unnormalized bounding boxes. auxiliary_outputs (`list[Dict]`, *optional*): Optional, only returned when auxiliary losses are activated (i.e. `config.auxiliary_loss` is set to `True`) and labels are provided. It is a list of dictionaries containing the two above keys (`logits` and `pred_boxes`) for each decoder layer. last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Sequence of hidden-states at the output of the last layer of the decoder of the model. """ loss: torch.FloatTensor | None = None loss_dict: dict | None = None logits: torch.FloatTensor | None = None pred_boxes: torch.FloatTensor | None = None auxiliary_outputs: list[dict] | None = None last_hidden_state: torch.FloatTensor | None = None decoder_hidden_states: tuple[torch.FloatTensor] | None = None decoder_attentions: tuple[torch.FloatTensor] | None = None cross_attentions: tuple[torch.FloatTensor] | None = None encoder_last_hidden_state: torch.FloatTensor | None = None encoder_hidden_states: tuple[torch.FloatTensor] | None = None encoder_attentions: tuple[torch.FloatTensor] | None = None @dataclass @auto_docstring( custom_intro=""" Output type of [`ConditionalDetrForSegmentation`]. """ ) class ConditionalDetrSegmentationOutput(ModelOutput): r""" loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` are provided)): Total loss as a linear combination of a negative log-likehood (cross-entropy) for class prediction and a bounding box loss. The latter is defined as a linear combination of the L1 loss and the generalized scale-invariant IoU loss. loss_dict (`Dict`, *optional*): A dictionary containing the individual losses. Useful for logging. logits (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes + 1)`): Classification logits (including no-object) for all queries. pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`): Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These values are normalized in [0, 1], relative to the size of each individual image in the batch (disregarding possible padding). You can use [`~ConditionalDetrImageProcessor.post_process_object_detection`] to retrieve the unnormalized bounding boxes. pred_masks (`torch.FloatTensor` of shape `(batch_size, num_queries, height/4, width/4)`): Segmentation masks logits for all queries. See also [`~ConditionalDetrImageProcessor.post_process_semantic_segmentation`] or [`~ConditionalDetrImageProcessor.post_process_instance_segmentation`] [`~ConditionalDetrImageProcessor.post_process_panoptic_segmentation`] to evaluate semantic, instance and panoptic segmentation masks respectively. auxiliary_outputs (`list[Dict]`, *optional*): Optional, only returned when auxiliary losses are activated (i.e. `config.auxiliary_loss` is set to `True`) and labels are provided. It is a list of dictionaries containing the two above keys (`logits` and `pred_boxes`) for each decoder layer. last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Sequence of hidden-states at the output of the last layer of the decoder of the model. """ loss: torch.FloatTensor | None = None loss_dict: dict | None = None logits: torch.FloatTensor | None = None pred_boxes: torch.FloatTensor | None = None pred_masks: torch.FloatTensor | None = None auxiliary_outputs: list[dict] | None = None last_hidden_state: torch.FloatTensor | None = None decoder_hidden_states: tuple[torch.FloatTensor] | None = None decoder_attentions: tuple[torch.FloatTensor] | None = None cross_attentions: tuple[torch.FloatTensor] | None = None encoder_last_hidden_state: torch.FloatTensor | None = None encoder_hidden_states: tuple[torch.FloatTensor] | None = None encoder_attentions: tuple[torch.FloatTensor] | None = None class ConditionalDetrFrozenBatchNorm2d(nn.Module): """ BatchNorm2d where the batch statistics and the affine parameters are fixed. Copy-paste from torchvision.misc.ops with added eps before rqsrt, without which any other models than torchvision.models.resnet[18,34,50,101] produce nans. """ def __init__(self, n): super().__init__() self.register_buffer("weight", torch.ones(n)) self.register_buffer("bias", torch.zeros(n)) self.register_buffer("running_mean", torch.zeros(n)) self.register_buffer("running_var", torch.ones(n)) def _load_from_state_dict( self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs ): num_batches_tracked_key = prefix + "num_batches_tracked" if num_batches_tracked_key in state_dict: del state_dict[num_batches_tracked_key] super()._load_from_state_dict( state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs ) def forward(self, x): # move reshapes to the beginning # to make it user-friendly weight = self.weight.reshape(1, -1, 1, 1) bias = self.bias.reshape(1, -1, 1, 1) running_var = self.running_var.reshape(1, -1, 1, 1) running_mean = self.running_mean.reshape(1, -1, 1, 1) epsilon = 1e-5 scale = weight * (running_var + epsilon).rsqrt() bias = bias - running_mean * scale return x * scale + bias def replace_batch_norm(model): r""" Recursively replace all `torch.nn.BatchNorm2d` with `ConditionalDetrFrozenBatchNorm2d`. Args: model (torch.nn.Module): input model """ for name, module in model.named_children(): if isinstance(module, nn.BatchNorm2d): new_module = ConditionalDetrFrozenBatchNorm2d(module.num_features) if module.weight.device != torch.device("meta"): new_module.weight.copy_(module.weight) new_module.bias.copy_(module.bias) new_module.running_mean.copy_(module.running_mean) new_module.running_var.copy_(module.running_var) model._modules[name] = new_module if len(list(module.children())) > 0: replace_batch_norm(module) class ConditionalDetrConvEncoder(nn.Module): """ Convolutional backbone, using either the AutoBackbone API or one from the timm library. nn.BatchNorm2d layers are replaced by ConditionalDetrFrozenBatchNorm2d as defined above. """ def __init__(self, config): super().__init__() self.config = config backbone = load_backbone(config) self.intermediate_channel_sizes = backbone.channels # replace batch norm by frozen batch norm with torch.no_grad(): replace_batch_norm(backbone) # We used to load with timm library directly instead of the AutoBackbone API # so we need to unwrap the `backbone._backbone` module to load weights without mismatch is_timm_model = False if hasattr(backbone, "_backbone"): backbone = backbone._backbone is_timm_model = True self.model = backbone backbone_model_type = config.backbone_config.model_type if "resnet" in backbone_model_type: for name, parameter in self.model.named_parameters(): if is_timm_model: if "layer2" not in name and "layer3" not in name and "layer4" not in name: parameter.requires_grad_(False) else: if "stage.1" not in name and "stage.2" not in name and "stage.3" not in name: parameter.requires_grad_(False) def forward(self, pixel_values: torch.Tensor, pixel_mask: torch.Tensor): # send pixel_values through the model to get list of feature maps features = self.model(pixel_values) if isinstance(features, dict): features = features.feature_maps out = [] for feature_map in features: # downsample pixel_mask to match shape of corresponding feature_map mask = nn.functional.interpolate(pixel_mask[None].float(), size=feature_map.shape[-2:]).to(torch.bool)[0] out.append((feature_map, mask)) return out class ConditionalDetrSinePositionEmbedding(nn.Module): """ This is a more standard version of the position embedding, very similar to the one used by the Attention is all you need paper, generalized to work on images. """ def __init__( self, num_position_features: int = 64, temperature: int = 10000, normalize: bool = False, scale: float | None = None, ): super().__init__() if scale is not None and normalize is False: raise ValueError("normalize should be True if scale is passed") self.num_position_features = num_position_features self.temperature = temperature self.normalize = normalize self.scale = 2 * math.pi if scale is None else scale @compile_compatible_method_lru_cache(maxsize=1) def forward( self, shape: torch.Size, device: torch.device | str, dtype: torch.dtype, mask: torch.Tensor | None = None, ) -> torch.Tensor: if mask is None: mask = torch.zeros((shape[0], shape[2], shape[3]), device=device, dtype=torch.bool) y_embed = mask.cumsum(1, dtype=dtype) x_embed = mask.cumsum(2, dtype=dtype) if self.normalize: eps = 1e-6 y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale dim_t = torch.arange(self.num_position_features, dtype=torch.int64, device=device).to(dtype) dim_t = self.temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / self.num_position_features) pos_x = x_embed[:, :, :, None] / dim_t pos_y = y_embed[:, :, :, None] / dim_t pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) # Flatten spatial dimensions and permute to (batch_size, sequence_length, hidden_size) format # expected by the encoder pos = pos.flatten(2).permute(0, 2, 1) return pos class ConditionalDetrLearnedPositionEmbedding(nn.Module): """ This module learns positional embeddings up to a fixed maximum size. """ def __init__(self, embedding_dim=256): super().__init__() self.row_embeddings = nn.Embedding(50, embedding_dim) self.column_embeddings = nn.Embedding(50, embedding_dim) @compile_compatible_method_lru_cache(maxsize=1) def forward( self, shape: torch.Size, device: torch.device | str, dtype: torch.dtype, mask: torch.Tensor | None = None, ): height, width = shape[-2:] width_values = torch.arange(width, device=device) height_values = torch.arange(height, device=device) x_emb = self.column_embeddings(width_values) y_emb = self.row_embeddings(height_values) pos = torch.cat([x_emb.unsqueeze(0).repeat(height, 1, 1), y_emb.unsqueeze(1).repeat(1, width, 1)], dim=-1) pos = pos.permute(2, 0, 1) pos = pos.unsqueeze(0) pos = pos.repeat(shape[0], 1, 1, 1) # Flatten spatial dimensions and permute to (batch_size, sequence_length, hidden_size) format # expected by the encoder pos = pos.flatten(2).permute(0, 2, 1) return pos def eager_attention_forward( module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: torch.Tensor | None, scaling: float | None = None, dropout: float = 0.0, **kwargs: Unpack[TransformersKwargs], ): if scaling is None: scaling = query.size(-1) ** -0.5 # Take the dot product between "query" and "key" to get the raw attention scores. attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling if attention_mask is not None: attn_weights = attn_weights + attention_mask attn_weights = nn.functional.softmax(attn_weights, dim=-1) attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) attn_output = torch.matmul(attn_weights, value) attn_output = attn_output.transpose(1, 2).contiguous() return attn_output, attn_weights class ConditionalDetrSelfAttention(nn.Module): """ Multi-headed self-attention from 'Attention Is All You Need' paper. In CONDITIONAL_DETR, position embeddings are added to both queries and keys (but not values) in self-attention. """ def __init__( self, config: ConditionalDetrConfig, hidden_size: int, num_attention_heads: int, dropout: float = 0.0, bias: bool = True, ): super().__init__() self.config = config self.head_dim = hidden_size // num_attention_heads self.scaling = self.head_dim**-0.5 self.attention_dropout = dropout self.is_causal = False self.k_proj = nn.Linear(hidden_size, hidden_size, bias=bias) self.v_proj = nn.Linear(hidden_size, hidden_size, bias=bias) self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias) self.o_proj = nn.Linear(hidden_size, hidden_size, bias=bias) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, position_embeddings: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor, torch.Tensor]: """ Position embeddings are added to both queries and keys (but not values). """ input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) query_key_input = hidden_states + position_embeddings if position_embeddings is not None else hidden_states query_states = self.q_proj(query_key_input).view(hidden_shape).transpose(1, 2) key_states = self.k_proj(query_key_input).view(hidden_shape).transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward ) attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, **kwargs, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights class ConditionalDetrDecoderSelfAttention(nn.Module): """ Multi-headed self-attention for Conditional DETR decoder layers. This attention module handles separate content and position projections, which are then combined before applying standard self-attention. Position embeddings are added to both queries and keys. """ def __init__( self, config: ConditionalDetrConfig, hidden_size: int, num_attention_heads: int, dropout: float = 0.0, ): super().__init__() self.config = config self.hidden_size = hidden_size self.head_dim = hidden_size // num_attention_heads self.scaling = self.head_dim**-0.5 self.attention_dropout = dropout self.is_causal = False # Content and position projections self.q_content_proj = nn.Linear(hidden_size, hidden_size) self.q_pos_proj = nn.Linear(hidden_size, hidden_size) self.k_content_proj = nn.Linear(hidden_size, hidden_size) self.k_pos_proj = nn.Linear(hidden_size, hidden_size) self.v_proj = nn.Linear(hidden_size, hidden_size) self.o_proj = nn.Linear(hidden_size, hidden_size) def forward( self, hidden_states: torch.Tensor, query_position_embeddings: torch.Tensor, attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor, torch.Tensor]: """ Args: hidden_states (`torch.Tensor` of shape `(batch_size, num_queries, hidden_size)`): Input hidden states from the decoder layer. query_position_embeddings (`torch.Tensor` of shape `(batch_size, num_queries, hidden_size)`): Position embeddings for queries and keys. Required (unlike standard attention). Processed through separate position projections (`q_pos_proj`, `k_pos_proj`) and added to content projections. attention_mask (`torch.Tensor` of shape `(batch_size, 1, num_queries, num_queries)`, *optional*): Attention mask to avoid attending to padding tokens. """ input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) query_states = ( (self.q_content_proj(hidden_states) + self.q_pos_proj(query_position_embeddings)) .view(hidden_shape) .transpose(1, 2) ) key_states = ( (self.k_content_proj(hidden_states) + self.k_pos_proj(query_position_embeddings)) .view(hidden_shape) .transpose(1, 2) ) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward ) attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, **kwargs, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights class ConditionalDetrDecoderCrossAttention(nn.Module): """ Multi-headed cross-attention for Conditional DETR decoder layers. This attention module handles the special cross-attention logic in Conditional DETR: - Separate content and position projections for queries and keys - Concatenation of query sine embeddings with queries (doubling query dimension) - Concatenation of key position embeddings with keys (doubling key dimension) - Output dimension remains hidden_size despite doubled input dimensions """ def __init__( self, config: ConditionalDetrConfig, hidden_size: int, num_attention_heads: int, dropout: float = 0.0, ): super().__init__() self.config = config self.hidden_size = hidden_size self.num_attention_heads = num_attention_heads self.head_dim = hidden_size // num_attention_heads self.attention_dropout = dropout self.is_causal = False # Content and position projections self.q_content_proj = nn.Linear(hidden_size, hidden_size) self.q_pos_proj = nn.Linear(hidden_size, hidden_size) self.k_content_proj = nn.Linear(hidden_size, hidden_size) self.k_pos_proj = nn.Linear(hidden_size, hidden_size) self.v_proj = nn.Linear(hidden_size, hidden_size) self.q_pos_sine_proj = nn.Linear(hidden_size, hidden_size) # Output projection: input is hidden_size * 2 (from concatenated q/k), output is hidden_size self.o_proj = nn.Linear(hidden_size, hidden_size) # Compute scaling for expanded head_dim (q and k have doubled dimensions after concatenation) # This matches the original Conditional DETR implementation where embed_dim * 2 is used expanded_head_dim = (hidden_size * 2) // num_attention_heads self.scaling = expanded_head_dim**-0.5 def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, query_sine_embed: torch.Tensor, encoder_position_embeddings: torch.Tensor, query_position_embeddings: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor, torch.Tensor]: """ Args: hidden_states (`torch.Tensor` of shape `(batch_size, num_queries, hidden_size)`): Decoder hidden states (queries). encoder_hidden_states (`torch.Tensor` of shape `(batch_size, encoder_seq_len, hidden_size)`): Encoder output hidden states (keys and values). query_sine_embed (`torch.Tensor` of shape `(batch_size, num_queries, hidden_size)`): Sine position embeddings for queries. **Concatenated** (not added) with query content, doubling the query dimension. encoder_position_embeddings (`torch.Tensor` of shape `(batch_size, encoder_seq_len, hidden_size)`): Position embeddings for keys. **Concatenated** (not added) with key content, doubling the key dimension. query_position_embeddings (`torch.Tensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*): Additional position embeddings. When provided (first layer only), **added** to query content before concatenation with `query_sine_embed`. Also causes `encoder_position_embeddings` to be added to key content before concatenation. attention_mask (`torch.Tensor` of shape `(batch_size, 1, num_queries, encoder_seq_len)`, *optional*): Attention mask to avoid attending to padding tokens. """ query_input_shape = hidden_states.shape[:-1] kv_input_shape = encoder_hidden_states.shape[:-1] query_hidden_shape = (*query_input_shape, self.num_attention_heads, self.head_dim) kv_hidden_shape = (*kv_input_shape, self.num_attention_heads, self.head_dim) # Apply content and position projections query_input = self.q_content_proj(hidden_states) key_input = self.k_content_proj(encoder_hidden_states) value_states = self.v_proj(encoder_hidden_states) key_pos = self.k_pos_proj(encoder_position_embeddings) # Combine content and position embeddings if query_position_embeddings is not None: query_input = query_input + self.q_pos_proj(query_position_embeddings) key_input = key_input + key_pos # Reshape and concatenate position embeddings (doubling head_dim) query_input = query_input.view(query_hidden_shape) key_input = key_input.view(kv_hidden_shape) query_sine_embed = self.q_pos_sine_proj(query_sine_embed).view(query_hidden_shape) key_pos = key_pos.view(kv_hidden_shape) query_states = torch.cat([query_input, query_sine_embed], dim=-1).view(*query_input_shape, -1) key_states = torch.cat([key_input, key_pos], dim=-1).view(*kv_input_shape, -1) # Reshape for attention computation expanded_head_dim = query_states.shape[-1] // self.num_attention_heads query_states = query_states.view(*query_input_shape, self.num_attention_heads, expanded_head_dim).transpose( 1, 2 ) key_states = key_states.view(*kv_input_shape, self.num_attention_heads, expanded_head_dim).transpose(1, 2) value_states = value_states.view(kv_hidden_shape).transpose(1, 2) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward ) attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, **kwargs, ) attn_output = attn_output.reshape(*query_input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights class ConditionalDetrMLP(nn.Module): def __init__(self, config: ConditionalDetrConfig, hidden_size: int, intermediate_size: int): super().__init__() self.fc1 = nn.Linear(hidden_size, intermediate_size) self.fc2 = nn.Linear(intermediate_size, hidden_size) self.activation_fn = ACT2FN[config.activation_function] self.activation_dropout = config.activation_dropout self.dropout = config.dropout def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.activation_fn(self.fc1(hidden_states)) hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) hidden_states = self.fc2(hidden_states) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) return hidden_states class ConditionalDetrEncoderLayer(GradientCheckpointingLayer): def __init__(self, config: ConditionalDetrConfig): super().__init__() self.hidden_size = config.d_model self.self_attn = ConditionalDetrSelfAttention( config=config, hidden_size=self.hidden_size, num_attention_heads=config.encoder_attention_heads, dropout=config.attention_dropout, ) self.self_attn_layer_norm = nn.LayerNorm(self.hidden_size) self.dropout = config.dropout self.mlp = ConditionalDetrMLP(config, self.hidden_size, config.encoder_ffn_dim) self.final_layer_norm = nn.LayerNorm(self.hidden_size) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, spatial_position_embeddings: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: """ Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, hidden_size)` attention_mask (`torch.FloatTensor`): attention mask of size `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative values. spatial_position_embeddings (`torch.FloatTensor`, *optional*): Spatial position embeddings (2D positional encodings of image locations), to be added to both the queries and keys in self-attention (but not to values). """ residual = hidden_states hidden_states, _ = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_embeddings=spatial_position_embeddings, **kwargs, ) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) hidden_states = residual + hidden_states hidden_states = self.self_attn_layer_norm(hidden_states) residual = hidden_states hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states hidden_states = self.final_layer_norm(hidden_states) if self.training: if not torch.isfinite(hidden_states).all(): clamp_value = torch.finfo(hidden_states.dtype).max - 1000 hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) return hidden_states class ConditionalDetrDecoderLayer(GradientCheckpointingLayer): def __init__(self, config: ConditionalDetrConfig): super().__init__() self.hidden_size = config.d_model self.self_attn = ConditionalDetrDecoderSelfAttention( config=config, hidden_size=self.hidden_size, num_attention_heads=config.decoder_attention_heads, dropout=config.attention_dropout, ) self.dropout = config.dropout self.self_attn_layer_norm = nn.LayerNorm(self.hidden_size) self.encoder_attn = ConditionalDetrDecoderCrossAttention( config=config, hidden_size=self.hidden_size, num_attention_heads=config.decoder_attention_heads, dropout=config.attention_dropout, ) self.encoder_attn_layer_norm = nn.LayerNorm(self.hidden_size) self.mlp = ConditionalDetrMLP(config, self.hidden_size, config.decoder_ffn_dim) self.final_layer_norm = nn.LayerNorm(self.hidden_size) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, spatial_position_embeddings: torch.Tensor | None = None, query_position_embeddings: torch.Tensor | None = None, query_sine_embed: torch.Tensor | None = None, encoder_hidden_states: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, is_first: bool | None = False, **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: """ Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(seq_len, batch, embed_dim)` attention_mask (`torch.FloatTensor`): attention mask of size `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative values. spatial_position_embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Spatial position embeddings (2D positional encodings) that are added to the queries and keys in each self-attention layer. query_position_embeddings (`torch.FloatTensor`, *optional*): object_queries that are added to the queries and keys in the self-attention layer. encoder_hidden_states (`torch.FloatTensor`): cross attention input to the layer of shape `(seq_len, batch, embed_dim)` encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative values. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. """ residual = hidden_states hidden_states, _ = self.self_attn( hidden_states=hidden_states, query_position_embeddings=query_position_embeddings, attention_mask=attention_mask, **kwargs, ) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) hidden_states = residual + hidden_states hidden_states = self.self_attn_layer_norm(hidden_states) if encoder_hidden_states is not None: residual = hidden_states hidden_states, _ = self.encoder_attn( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=encoder_attention_mask, query_sine_embed=query_sine_embed, encoder_position_embeddings=spatial_position_embeddings, # Only pass query_position_embeddings for the first layer query_position_embeddings=query_position_embeddings if is_first else None, **kwargs, ) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) hidden_states = residual + hidden_states hidden_states = self.encoder_attn_layer_norm(hidden_states) # Fully Connected residual = hidden_states hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states hidden_states = self.final_layer_norm(hidden_states) return hidden_states class ConditionalDetrMLPPredictionHead(nn.Module): """ Very simple multi-layer perceptron (MLP, also called FFN), used to predict the normalized center coordinates, height and width of a bounding box w.r.t. an image. """ def __init__(self, input_dim, hidden_dim, output_dim, num_layers): super().__init__() self.num_layers = num_layers h = [hidden_dim] * (num_layers - 1) self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])) def forward(self, x): for i, layer in enumerate(self.layers): x = nn.functional.relu(layer(x)) if i < self.num_layers - 1 else layer(x) return x class ConditionalDetrConvBlock(nn.Module): """Basic conv block: Conv3x3 -> GroupNorm -> Activation.""" def __init__(self, in_channels: int, out_channels: int, activation: str = "relu"): super().__init__() self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) self.norm = nn.GroupNorm(min(8, out_channels), out_channels) self.activation = ACT2FN[activation] def forward(self, x: torch.Tensor) -> torch.Tensor: return self.activation(self.norm(self.conv(x))) class ConditionalDetrFPNFusionStage(nn.Module): """Single FPN fusion stage combining low-resolution features with high-resolution FPN features.""" def __init__(self, fpn_channels: int, current_channels: int, output_channels: int, activation: str = "relu"): super().__init__() self.fpn_adapter = nn.Conv2d(fpn_channels, current_channels, kernel_size=1) self.refine = ConditionalDetrConvBlock(current_channels, output_channels, activation) def forward(self, features: torch.Tensor, fpn_features: torch.Tensor) -> torch.Tensor: """ Args: features: Current features to upsample, shape (B*Q, current_channels, H_in, W_in) fpn_features: FPN features at target resolution, shape (B*Q, fpn_channels, H_out, W_out) Returns: Fused and refined features, shape (B*Q, output_channels, H_out, W_out) """ fpn_features = self.fpn_adapter(fpn_features) features = nn.functional.interpolate(features, size=fpn_features.shape[-2:], mode="nearest") return self.refine(fpn_features + features) class ConditionalDetrMaskHeadSmallConv(nn.Module): """ Segmentation mask head that generates per-query masks using FPN-based progressive upsampling. Combines attention maps (spatial localization) with encoder features (semantics) and progressively upsamples through multiple scales, fusing with FPN features for high-resolution detail. """ def __init__( self, input_channels: int, fpn_channels: list[int], hidden_size: int, activation_function: str = "relu", ): super().__init__() if input_channels % 8 != 0: raise ValueError(f"input_channels must be divisible by 8, got {input_channels}") self.conv1 = ConditionalDetrConvBlock(input_channels, input_channels, activation_function) self.conv2 = ConditionalDetrConvBlock(input_channels, hidden_size // 2, activation_function) # Progressive channel reduction: /2 -> /4 -> /8 -> /16 self.fpn_stages = nn.ModuleList( [ ConditionalDetrFPNFusionStage( fpn_channels[0], hidden_size // 2, hidden_size // 4, activation_function ), ConditionalDetrFPNFusionStage( fpn_channels[1], hidden_size // 4, hidden_size // 8, activation_function ), ConditionalDetrFPNFusionStage( fpn_channels[2], hidden_size // 8, hidden_size // 16, activation_function ), ] ) self.output_conv = nn.Conv2d(hidden_size // 16, 1, kernel_size=3, padding=1) def forward( self, features: torch.Tensor, attention_masks: torch.Tensor, fpn_features: list[torch.Tensor], ) -> torch.Tensor: """ Args: features: Encoder output features, shape (batch_size, hidden_size, H, W) attention_masks: Cross-attention maps from decoder, shape (batch_size, num_queries, num_heads, H, W) fpn_features: List of 3 FPN features from low to high resolution, each (batch_size, C, H, W) Returns: Predicted masks, shape (batch_size * num_queries, 1, output_H, output_W) """ num_queries = attention_masks.shape[1] # Expand to (batch_size * num_queries) dimension features = features.unsqueeze(1).expand(-1, num_queries, -1, -1, -1).flatten(0, 1) attention_masks = attention_masks.flatten(0, 1) fpn_features = [ fpn_feat.unsqueeze(1).expand(-1, num_queries, -1, -1, -1).flatten(0, 1) for fpn_feat in fpn_features ] hidden_states = torch.cat([features, attention_masks], dim=1) hidden_states = self.conv1(hidden_states) hidden_states = self.conv2(hidden_states) for fpn_stage, fpn_feat in zip(self.fpn_stages, fpn_features): hidden_states = fpn_stage(hidden_states, fpn_feat) return self.output_conv(hidden_states) class ConditionalDetrMHAttentionMap(nn.Module): """This is a 2D attention module, which only returns the attention softmax (no multiplication by value)""" def __init__( self, hidden_size: int, num_attention_heads: int, dropout: float = 0.0, bias: bool = True, ): super().__init__() self.head_dim = hidden_size // num_attention_heads self.scaling = self.head_dim**-0.5 self.attention_dropout = dropout self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias) self.k_proj = nn.Linear(hidden_size, hidden_size, bias=bias) def forward( self, query_states: torch.Tensor, key_states: torch.Tensor, attention_mask: torch.Tensor | None = None ): query_hidden_shape = (*query_states.shape[:-1], -1, self.head_dim) key_hidden_shape = (key_states.shape[0], -1, self.head_dim, *key_states.shape[-2:]) query_states = self.q_proj(query_states).view(query_hidden_shape) key_states = nn.functional.conv2d( key_states, self.k_proj.weight.unsqueeze(-1).unsqueeze(-1), self.k_proj.bias ).view(key_hidden_shape) batch_size, num_queries, num_heads, head_dim = query_states.shape _, _, _, height, width = key_states.shape query_shape = (batch_size * num_heads, num_queries, head_dim) key_shape = (batch_size * num_heads, height * width, head_dim) attn_weights_shape = (batch_size, num_heads, num_queries, height, width) query = query_states.transpose(1, 2).contiguous().view(query_shape) key = key_states.permute(0, 1, 3, 4, 2).contiguous().view(key_shape) attn_weights = ( (torch.matmul(query * self.scaling, key.transpose(1, 2))).view(attn_weights_shape).transpose(1, 2) ) if attention_mask is not None: attn_weights = attn_weights + attention_mask attn_weights = nn.functional.softmax(attn_weights.flatten(2), dim=-1).view(attn_weights.size()) attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) return attn_weights @auto_docstring class ConditionalDetrPreTrainedModel(PreTrainedModel): config: ConditionalDetrConfig base_model_prefix = "model" main_input_name = "pixel_values" input_modalities = ("image",) _no_split_modules = [r"ConditionalDetrConvEncoder", r"ConditionalDetrEncoderLayer", r"ConditionalDetrDecoderLayer"] supports_gradient_checkpointing = True _supports_sdpa = True _supports_flash_attn = True _supports_attention_backend = True _supports_flex_attn = True # Uses create_bidirectional_masks for attention masking _keys_to_ignore_on_load_unexpected = [ r"detr\.model\.backbone\.model\.layer\d+\.0\.downsample\.1\.num_batches_tracked" ] @torch.no_grad() def _init_weights(self, module): std = self.config.init_std xavier_std = self.config.init_xavier_std if isinstance(module, ConditionalDetrMaskHeadSmallConv): # ConditionalDetrMaskHeadSmallConv uses kaiming initialization for all its Conv2d layers for m in module.modules(): if isinstance(m, nn.Conv2d): init.kaiming_uniform_(m.weight, a=1) if m.bias is not None: init.constant_(m.bias, 0) elif isinstance(module, ConditionalDetrMHAttentionMap): init.zeros_(module.k_proj.bias) init.zeros_(module.q_proj.bias) init.xavier_uniform_(module.k_proj.weight, gain=xavier_std) init.xavier_uniform_(module.q_proj.weight, gain=xavier_std) elif isinstance(module, ConditionalDetrLearnedPositionEmbedding): init.uniform_(module.row_embeddings.weight) init.uniform_(module.column_embeddings.weight) elif isinstance(module, (nn.Linear, nn.Conv2d)): init.normal_(module.weight, mean=0.0, std=std) if module.bias is not None: init.zeros_(module.bias) elif isinstance(module, nn.Embedding): init.normal_(module.weight, mean=0.0, std=std) # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False): init.zeros_(module.weight[module.padding_idx]) elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)): init.ones_(module.weight) init.zeros_(module.bias) class ConditionalDetrEncoder(ConditionalDetrPreTrainedModel): """ Transformer encoder that processes a flattened feature map from a vision backbone, composed of a stack of [`ConditionalDetrEncoderLayer`] modules. Args: config (`ConditionalDetrConfig`): Model configuration object. """ _can_record_outputs = {"hidden_states": ConditionalDetrEncoderLayer, "attentions": ConditionalDetrSelfAttention} def __init__(self, config: ConditionalDetrConfig): super().__init__(config) self.dropout = config.dropout self.layers = nn.ModuleList([ConditionalDetrEncoderLayer(config) for _ in range(config.encoder_layers)]) # Initialize weights and apply final processing self.post_init() @merge_with_config_defaults @capture_outputs def forward( self, inputs_embeds=None, attention_mask=None, spatial_position_embeddings=None, **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutput: r""" Args: inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Flattened feature map (output of the backbone + projection layer) that is passed to the encoder. attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding pixel features. Mask values selected in `[0, 1]`: - 1 for pixel features that are real (i.e. **not masked**), - 0 for pixel features that are padding (i.e. **masked**). [What are attention masks?](../glossary#attention-mask) spatial_position_embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Spatial position embeddings (2D positional encodings) that are added to the queries and keys in each self-attention layer. """ hidden_states = inputs_embeds hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) attention_mask = create_bidirectional_mask( config=self.config, inputs_embeds=inputs_embeds, attention_mask=attention_mask, ) for encoder_layer in self.layers: # we add spatial_position_embeddings as extra input to the encoder_layer hidden_states = encoder_layer( hidden_states, attention_mask, spatial_position_embeddings=spatial_position_embeddings, **kwargs ) return BaseModelOutput(last_hidden_state=hidden_states) # function to generate sine positional embedding for 2d coordinates def gen_sine_position_embeddings(pos_tensor, d_model): scale = 2 * math.pi dim = d_model // 2 dim_t = torch.arange(dim, dtype=torch.float32, device=pos_tensor.device) dim_t = 10000 ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / dim) x_embed = pos_tensor[:, :, 0] * scale y_embed = pos_tensor[:, :, 1] * scale pos_x = x_embed[:, :, None] / dim_t pos_y = y_embed[:, :, None] / dim_t pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2) pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3).flatten(2) pos = torch.cat((pos_y, pos_x), dim=2) return pos.to(pos_tensor.dtype) class ConditionalDetrDecoder(ConditionalDetrPreTrainedModel): """ Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`ConditionalDetrDecoderLayer`]. The decoder updates the query embeddings through multiple self-attention and cross-attention layers. Some small tweaks for Conditional DETR: - object_queries and query_position_embeddings are added to the forward pass. - if self.config.auxiliary_loss is set to True, also returns a stack of activations from all decoding layers. Args: config: ConditionalDetrConfig """ _can_record_outputs = { "hidden_states": ConditionalDetrDecoderLayer, "attentions": OutputRecorder(ConditionalDetrDecoderSelfAttention, layer_name="self_attn", index=1), "cross_attentions": OutputRecorder(ConditionalDetrDecoderCrossAttention, layer_name="encoder_attn", index=1), } def __init__(self, config: ConditionalDetrConfig): super().__init__(config) self.hidden_size = config.d_model self.dropout = config.dropout self.layerdrop = config.decoder_layerdrop self.layers = nn.ModuleList([ConditionalDetrDecoderLayer(config) for _ in range(config.decoder_layers)]) # in Conditional DETR, the decoder uses layernorm after the last decoder layer output self.layernorm = nn.LayerNorm(config.d_model) # query_scale is the FFN applied on f to generate transformation T self.query_scale = ConditionalDetrMLPPredictionHead(self.hidden_size, self.hidden_size, self.hidden_size, 2) self.ref_point_head = ConditionalDetrMLPPredictionHead(self.hidden_size, self.hidden_size, 2, 2) for layer_id in range(config.decoder_layers - 1): # Set q_pos_proj to None for layers after the first (only first layer uses query position embeddings) self.layers[layer_id + 1].encoder_attn.q_pos_proj = None # Initialize weights and apply final processing self.post_init() @merge_with_config_defaults @capture_outputs def forward( self, inputs_embeds=None, attention_mask=None, encoder_hidden_states=None, encoder_attention_mask=None, spatial_position_embeddings=None, object_queries_position_embeddings=None, **kwargs: Unpack[TransformersKwargs], ) -> ConditionalDetrDecoderOutput: r""" Args: inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): The query embeddings that are passed into the decoder. attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on certain queries. Mask values selected in `[0, 1]`: - 1 for queries that are **not masked**, - 0 for queries that are **masked**. [What are attention masks?](../glossary#attention-mask) encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*): Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder. encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*): Mask to avoid performing cross-attention on padding pixel_values of the encoder. Mask values selected in `[0, 1]`: - 1 for pixels that are real (i.e. **not masked**), - 0 for pixels that are padding (i.e. **masked**). spatial_position_embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Spatial position embeddings that are added to the queries and keys in each cross-attention layer. object_queries_position_embeddings (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`): , *optional*): Position embeddings that are added to the queries and keys in each self-attention layer. """ if inputs_embeds is not None: hidden_states = inputs_embeds # expand encoder attention mask if encoder_hidden_states is not None and encoder_attention_mask is not None: # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len] encoder_attention_mask = create_bidirectional_mask( self.config, inputs_embeds, encoder_attention_mask, ) # optional intermediate hidden states intermediate = () if self.config.auxiliary_loss else None reference_points_before_sigmoid = self.ref_point_head( object_queries_position_embeddings ) # [num_queries, batch_size, 2] reference_points = reference_points_before_sigmoid.sigmoid().transpose(0, 1) obj_center = reference_points[..., :2].transpose(0, 1) # get sine embedding for the query vector query_sine_embed_before_transformation = gen_sine_position_embeddings(obj_center, self.config.d_model) for idx, decoder_layer in enumerate(self.layers): if self.training: dropout_probability = torch.rand([]) if dropout_probability < self.layerdrop: continue if idx == 0: pos_transformation = 1 else: pos_transformation = self.query_scale(hidden_states) # apply transformation query_sine_embed = query_sine_embed_before_transformation * pos_transformation hidden_states = decoder_layer( hidden_states, None, spatial_position_embeddings, object_queries_position_embeddings, query_sine_embed, encoder_hidden_states, # as a positional argument for gradient checkpointing encoder_attention_mask=encoder_attention_mask, is_first=(idx == 0), **kwargs, ) if self.config.auxiliary_loss: hidden_states = self.layernorm(hidden_states) intermediate += (hidden_states,) # finally, apply layernorm hidden_states = self.layernorm(hidden_states) # stack intermediate decoder activations if self.config.auxiliary_loss: intermediate = torch.stack(intermediate) return ConditionalDetrDecoderOutput( last_hidden_state=hidden_states, intermediate_hidden_states=intermediate, reference_points=reference_points, ) @auto_docstring( custom_intro=""" The bare CONDITIONAL_DETR Model (consisting of a backbone and encoder-decoder Transformer) outputting raw hidden-states without any specific head on top. """ ) class ConditionalDetrModel(ConditionalDetrPreTrainedModel): def __init__(self, config: ConditionalDetrConfig): super().__init__(config) self.backbone = ConditionalDetrConvEncoder(config) if config.position_embedding_type == "sine": self.position_embedding = ConditionalDetrSinePositionEmbedding(config.d_model // 2, normalize=True) elif config.position_embedding_type == "learned": self.position_embedding = ConditionalDetrLearnedPositionEmbedding(config.d_model // 2) else: raise ValueError(f"Not supported {config.position_embedding_type}") self.query_position_embeddings = nn.Embedding(config.num_queries, config.d_model) self.input_projection = nn.Conv2d(self.backbone.intermediate_channel_sizes[-1], config.d_model, kernel_size=1) self.encoder = ConditionalDetrEncoder(config) self.decoder = ConditionalDetrDecoder(config) # Initialize weights and apply final processing self.post_init() def freeze_backbone(self): for _, param in self.backbone.model.named_parameters(): param.requires_grad_(False) def unfreeze_backbone(self): for _, param in self.backbone.model.named_parameters(): param.requires_grad_(True) @auto_docstring @can_return_tuple def forward( self, pixel_values: torch.FloatTensor, pixel_mask: torch.LongTensor | None = None, decoder_attention_mask: torch.LongTensor | None = None, encoder_outputs: torch.FloatTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, decoder_inputs_embeds: torch.FloatTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> ConditionalDetrModelOutput: r""" decoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, num_queries)`, *optional*): Not used by default. Can be used to mask object queries. inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing the flattened feature map (output of the backbone + projection layer), you can choose to directly pass a flattened representation of an image. decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*): Optionally, instead of initializing the queries with a tensor of zeros, you can choose to directly pass an embedded representation. Examples: ```python >>> from transformers import AutoImageProcessor, AutoModel >>> from PIL import Image >>> import requests >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/conditional-detr-resnet-50") >>> model = AutoModel.from_pretrained("microsoft/conditional-detr-resnet-50") >>> # prepare image for the model >>> inputs = image_processor(images=image, return_tensors="pt") >>> # forward pass >>> outputs = model(**inputs) >>> # the last hidden states are the final query embeddings of the Transformer decoder >>> # these are of shape (batch_size, num_queries, hidden_size) >>> last_hidden_states = outputs.last_hidden_state >>> list(last_hidden_states.shape) [1, 300, 256] ```""" batch_size, num_channels, height, width = pixel_values.shape device = pixel_values.device if pixel_mask is None: pixel_mask = torch.ones(((batch_size, height, width)), device=device) # First, sent pixel_values + pixel_mask through Backbone to obtain the features # pixel_values should be of shape (batch_size, num_channels, height, width) # pixel_mask should be of shape (batch_size, height, width) features = self.backbone(pixel_values, pixel_mask) # get final feature map and downsampled mask feature_map, mask = features[-1] if mask is None: raise ValueError("Backbone does not return downsampled pixel mask") # Second, apply 1x1 convolution to reduce the channel dimension to d_model (256 by default) projected_feature_map = self.input_projection(feature_map) # Generate position embeddings spatial_position_embeddings = self.position_embedding( shape=feature_map.shape, device=device, dtype=pixel_values.dtype, mask=mask ) # Third, flatten the feature map of shape NxCxHxW to NxCxHW, and permute it to NxHWxC # In other words, turn their shape into (batch_size, sequence_length, hidden_size) flattened_features = projected_feature_map.flatten(2).permute(0, 2, 1) flattened_mask = mask.flatten(1) # Fourth, sent flattened_features + flattened_mask + spatial_position_embeddings through encoder # flattened_features is a Tensor of shape (batch_size, height*width, hidden_size) # flattened_mask is a Tensor of shape (batch_size, height*width) if encoder_outputs is None: encoder_outputs = self.encoder( inputs_embeds=flattened_features, attention_mask=flattened_mask, spatial_position_embeddings=spatial_position_embeddings, **kwargs, ) # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput elif not isinstance(encoder_outputs, BaseModelOutput): encoder_outputs = BaseModelOutput( last_hidden_state=encoder_outputs[0], hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, ) # Fifth, sent query embeddings through the decoder (which is conditioned on the encoder output) object_queries_position_embeddings = self.query_position_embeddings.weight.unsqueeze(0).repeat( batch_size, 1, 1 ) queries = torch.zeros_like(object_queries_position_embeddings) # decoder outputs consists of (dec_features, dec_hidden, dec_attn) decoder_outputs = self.decoder( inputs_embeds=queries, attention_mask=None, spatial_position_embeddings=spatial_position_embeddings, object_queries_position_embeddings=object_queries_position_embeddings, encoder_hidden_states=encoder_outputs.last_hidden_state, encoder_attention_mask=flattened_mask, **kwargs, ) return ConditionalDetrModelOutput( last_hidden_state=decoder_outputs.last_hidden_state, decoder_hidden_states=decoder_outputs.hidden_states, decoder_attentions=decoder_outputs.attentions, cross_attentions=decoder_outputs.cross_attentions, encoder_last_hidden_state=encoder_outputs.last_hidden_state, encoder_hidden_states=encoder_outputs.hidden_states, encoder_attentions=encoder_outputs.attentions, intermediate_hidden_states=decoder_outputs.intermediate_hidden_states, reference_points=decoder_outputs.reference_points, ) def inverse_sigmoid(x, eps=1e-5): x = x.clamp(min=0, max=1) x1 = x.clamp(min=eps) x2 = (1 - x).clamp(min=eps) return torch.log(x1 / x2) @auto_docstring( custom_intro=""" CONDITIONAL_DETR Model (consisting of a backbone and encoder-decoder Transformer) with object detection heads on top, for tasks such as COCO detection. """ ) class ConditionalDetrForObjectDetection(ConditionalDetrPreTrainedModel): def __init__(self, config: ConditionalDetrConfig): super().__init__(config) # CONDITIONAL_DETR encoder-decoder model self.model = ConditionalDetrModel(config) self.class_labels_classifier = nn.Linear(config.d_model, config.num_labels) self.bbox_predictor = ConditionalDetrMLPPredictionHead( input_dim=config.d_model, hidden_dim=config.d_model, output_dim=4, num_layers=3 ) # Initialize weights and apply final processing self.post_init() @auto_docstring @can_return_tuple def forward( self, pixel_values: torch.FloatTensor, pixel_mask: torch.LongTensor | None = None, decoder_attention_mask: torch.LongTensor | None = None, encoder_outputs: torch.FloatTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, decoder_inputs_embeds: torch.FloatTensor | None = None, labels: list[dict] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> ConditionalDetrObjectDetectionOutput: r""" decoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, num_queries)`, *optional*): Not used by default. Can be used to mask object queries. inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing the flattened feature map (output of the backbone + projection layer), you can choose to directly pass a flattened representation of an image. decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*): Optionally, instead of initializing the queries with a tensor of zeros, you can choose to directly pass an embedded representation. labels (`list[Dict]` of len `(batch_size,)`, *optional*): Labels for computing the bipartite matching loss. List of dicts, each dictionary containing at least the following 2 keys: 'class_labels' and 'boxes' (the class labels and bounding boxes of an image in the batch respectively). The class labels themselves should be a `torch.LongTensor` of len `(number of bounding boxes in the image,)` and the boxes a `torch.FloatTensor` of shape `(number of bounding boxes in the image, 4)`. Examples: ```python >>> from transformers import AutoImageProcessor, AutoModelForObjectDetection >>> from PIL import Image >>> import requests >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/conditional-detr-resnet-50") >>> model = AutoModelForObjectDetection.from_pretrained("microsoft/conditional-detr-resnet-50") >>> inputs = image_processor(images=image, return_tensors="pt") >>> outputs = model(**inputs) >>> # convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax) >>> target_sizes = torch.tensor([image.size[::-1]]) >>> results = image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[ ... 0 ... ] >>> for score, label, box in zip(results["scores"], results["labels"], results["boxes"]): ... box = [round(i, 2) for i in box.tolist()] ... print( ... f"Detected {model.config.id2label[label.item()]} with confidence " ... f"{round(score.item(), 3)} at location {box}" ... ) Detected remote with confidence 0.833 at location [38.31, 72.1, 177.63, 118.45] Detected cat with confidence 0.831 at location [9.2, 51.38, 321.13, 469.0] Detected cat with confidence 0.804 at location [340.3, 16.85, 642.93, 370.95] Detected remote with confidence 0.683 at location [334.48, 73.49, 366.37, 190.01] Detected couch with confidence 0.535 at location [0.52, 1.19, 640.35, 475.1] ```""" # First, sent images through CONDITIONAL_DETR base model to obtain encoder + decoder outputs outputs = self.model( pixel_values, pixel_mask=pixel_mask, decoder_attention_mask=decoder_attention_mask, encoder_outputs=encoder_outputs, inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, **kwargs, ) sequence_output = outputs[0] # class logits + predicted bounding boxes logits = self.class_labels_classifier(sequence_output) reference = outputs.reference_points reference_before_sigmoid = inverse_sigmoid(reference).transpose(0, 1) hs = sequence_output tmp = self.bbox_predictor(hs) tmp[..., :2] += reference_before_sigmoid pred_boxes = tmp.sigmoid() # pred_boxes = self.bbox_predictor(sequence_output).sigmoid() loss, loss_dict, auxiliary_outputs = None, None, None if labels is not None: outputs_class, outputs_coord = None, None if self.config.auxiliary_loss: outputs_coords = [] intermediate = outputs.intermediate_hidden_states outputs_class = self.class_labels_classifier(intermediate) for lvl in range(intermediate.shape[0]): tmp = self.bbox_predictor(intermediate[lvl]) tmp[..., :2] += reference_before_sigmoid outputs_coord = tmp.sigmoid() outputs_coords.append(outputs_coord) outputs_coord = torch.stack(outputs_coords) loss, loss_dict, auxiliary_outputs = self.loss_function( logits, labels, self.device, pred_boxes, self.config, outputs_class, outputs_coord ) return ConditionalDetrObjectDetectionOutput( loss=loss, loss_dict=loss_dict, logits=logits, pred_boxes=pred_boxes, auxiliary_outputs=auxiliary_outputs, last_hidden_state=outputs.last_hidden_state, decoder_hidden_states=outputs.decoder_hidden_states, decoder_attentions=outputs.decoder_attentions, cross_attentions=outputs.cross_attentions, encoder_last_hidden_state=outputs.encoder_last_hidden_state, encoder_hidden_states=outputs.encoder_hidden_states, encoder_attentions=outputs.encoder_attentions, ) # taken from https://github.com/Atten4Vis/conditionalDETR/blob/master/models/conditional_detr.py def _set_aux_loss(self, outputs_class, outputs_coord): return [{"logits": a, "pred_boxes": b} for a, b in zip(outputs_class[:-1], outputs_coord[:-1])] @auto_docstring( custom_intro=""" CONDITIONAL_DETR Model (consisting of a backbone and encoder-decoder Transformer) with a segmentation head on top, for tasks such as COCO panoptic. """ ) class ConditionalDetrForSegmentation(ConditionalDetrPreTrainedModel): _checkpoint_conversion_mapping = { "bbox_attention.q_linear": "bbox_attention.q_proj", "bbox_attention.k_linear": "bbox_attention.k_proj", # Mask head refactor "mask_head.lay1": "mask_head.conv1.conv", "mask_head.gn1": "mask_head.conv1.norm", "mask_head.lay2": "mask_head.conv2.conv", "mask_head.gn2": "mask_head.conv2.norm", "mask_head.adapter1": "mask_head.fpn_stages.0.fpn_adapter", "mask_head.lay3": "mask_head.fpn_stages.0.refine.conv", "mask_head.gn3": "mask_head.fpn_stages.0.refine.norm", "mask_head.adapter2": "mask_head.fpn_stages.1.fpn_adapter", "mask_head.lay4": "mask_head.fpn_stages.1.refine.conv", "mask_head.gn4": "mask_head.fpn_stages.1.refine.norm", "mask_head.adapter3": "mask_head.fpn_stages.2.fpn_adapter", "mask_head.lay5": "mask_head.fpn_stages.2.refine.conv", "mask_head.gn5": "mask_head.fpn_stages.2.refine.norm", "mask_head.out_lay": "mask_head.output_conv", } def __init__(self, config: ConditionalDetrConfig): super().__init__(config) # object detection model self.conditional_detr = ConditionalDetrForObjectDetection(config) # segmentation head hidden_size, number_of_heads = config.d_model, config.encoder_attention_heads intermediate_channel_sizes = self.conditional_detr.model.backbone.intermediate_channel_sizes self.mask_head = ConditionalDetrMaskHeadSmallConv( input_channels=hidden_size + number_of_heads, fpn_channels=intermediate_channel_sizes[::-1][-3:], hidden_size=hidden_size, activation_function=config.activation_function, ) self.bbox_attention = ConditionalDetrMHAttentionMap(hidden_size, number_of_heads, dropout=0.0) # Initialize weights and apply final processing self.post_init() @auto_docstring @can_return_tuple def forward( self, pixel_values: torch.FloatTensor, pixel_mask: torch.LongTensor | None = None, decoder_attention_mask: torch.FloatTensor | None = None, encoder_outputs: torch.FloatTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, decoder_inputs_embeds: torch.FloatTensor | None = None, labels: list[dict] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor] | ConditionalDetrSegmentationOutput: r""" decoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, num_queries)`, *optional*): Mask to avoid performing attention on certain object queries in the decoder. Mask values selected in `[0, 1]`: - 1 for queries that are **not masked**, - 0 for queries that are **masked**. inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Kept for backward compatibility, but cannot be used for segmentation, as segmentation requires multi-scale features from the backbone that are not available when bypassing it with inputs_embeds. decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*): Optionally, instead of initializing the queries with a tensor of zeros, you can choose to directly pass an embedded representation. Useful for tasks that require custom query initialization. labels (`list[Dict]` of len `(batch_size,)`, *optional*): Labels for computing the bipartite matching loss, DICE/F-1 loss and Focal loss. List of dicts, each dictionary containing at least the following 3 keys: 'class_labels', 'boxes' and 'masks' (the class labels, bounding boxes and segmentation masks of an image in the batch respectively). The class labels themselves should be a `torch.LongTensor` of len `(number of bounding boxes in the image,)`, the boxes a `torch.FloatTensor` of shape `(number of bounding boxes in the image, 4)` and the masks a `torch.FloatTensor` of shape `(number of bounding boxes in the image, height, width)`. Examples: ```python >>> import io >>> import httpx >>> from io import BytesIO >>> from PIL import Image >>> import torch >>> import numpy >>> from transformers import AutoImageProcessor, ConditionalDetrForSegmentation >>> from transformers.image_transforms import rgb_to_id >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> with httpx.stream("GET", url) as response: ... image = Image.open(BytesIO(response.read())) >>> image_processor = AutoImageProcessor.from_pretrained("facebook/conditional_detr-resnet-50-panoptic") >>> model = ConditionalDetrForSegmentation.from_pretrained("facebook/conditional_detr-resnet-50-panoptic") >>> # prepare image for the model >>> inputs = image_processor(images=image, return_tensors="pt") >>> # forward pass >>> outputs = model(**inputs) >>> # Use the `post_process_panoptic_segmentation` method of the `image_processor` to retrieve post-processed panoptic segmentation maps >>> # Segmentation results are returned as a list of dictionaries >>> result = image_processor.post_process_panoptic_segmentation(outputs, target_sizes=[(300, 500)]) >>> # A tensor of shape (height, width) where each value denotes a segment id, filled with -1 if no segment is found >>> panoptic_seg = result[0]["segmentation"] >>> panoptic_seg.shape torch.Size([300, 500]) >>> # Get prediction score and segment_id to class_id mapping of each segment >>> panoptic_segments_info = result[0]["segments_info"] >>> len(panoptic_segments_info) 5 ```""" batch_size, num_channels, height, width = pixel_values.shape device = pixel_values.device if pixel_mask is None: pixel_mask = torch.ones((batch_size, height, width), device=device) vision_features = self.conditional_detr.model.backbone(pixel_values, pixel_mask) feature_map, mask = vision_features[-1] # Apply 1x1 conv to map (batch_size, C, H, W) -> (batch_size, hidden_size, H, W), then flatten to (batch_size, HW, hidden_size) projected_feature_map = self.conditional_detr.model.input_projection(feature_map) flattened_features = projected_feature_map.flatten(2).permute(0, 2, 1) spatial_position_embeddings = self.conditional_detr.model.position_embedding( shape=feature_map.shape, device=device, dtype=pixel_values.dtype, mask=mask ) flattened_mask = mask.flatten(1) if encoder_outputs is None: encoder_outputs = self.conditional_detr.model.encoder( inputs_embeds=flattened_features, attention_mask=flattened_mask, spatial_position_embeddings=spatial_position_embeddings, **kwargs, ) object_queries_position_embeddings = self.conditional_detr.model.query_position_embeddings.weight.unsqueeze( 0 ).repeat(batch_size, 1, 1) # Use decoder_inputs_embeds as queries if provided, otherwise initialize with zeros if decoder_inputs_embeds is not None: queries = decoder_inputs_embeds else: queries = torch.zeros_like(object_queries_position_embeddings) decoder_outputs = self.conditional_detr.model.decoder( inputs_embeds=queries, attention_mask=decoder_attention_mask, spatial_position_embeddings=spatial_position_embeddings, object_queries_position_embeddings=object_queries_position_embeddings, encoder_hidden_states=encoder_outputs.last_hidden_state, encoder_attention_mask=flattened_mask, **kwargs, ) sequence_output = decoder_outputs[0] logits = self.conditional_detr.class_labels_classifier(sequence_output) pred_boxes = self.conditional_detr.bbox_predictor(sequence_output).sigmoid() height, width = feature_map.shape[-2:] memory = encoder_outputs.last_hidden_state.permute(0, 2, 1).view( batch_size, self.config.d_model, height, width ) attention_mask = flattened_mask.view(batch_size, height, width) if attention_mask is not None: min_dtype = torch.finfo(memory.dtype).min attention_mask = torch.where( attention_mask.unsqueeze(1).unsqueeze(1), torch.tensor(0.0, device=memory.device, dtype=memory.dtype), min_dtype, ) bbox_mask = self.bbox_attention(sequence_output, memory, attention_mask=attention_mask) seg_masks = self.mask_head( features=projected_feature_map, attention_masks=bbox_mask, fpn_features=[vision_features[2][0], vision_features[1][0], vision_features[0][0]], ) pred_masks = seg_masks.view( batch_size, self.conditional_detr.config.num_queries, seg_masks.shape[-2], seg_masks.shape[-1] ) loss, loss_dict, auxiliary_outputs = None, None, None if labels is not None: outputs_class, outputs_coord = None, None if self.config.auxiliary_loss: intermediate = decoder_outputs.intermediate_hidden_states outputs_class = self.conditional_detr.class_labels_classifier(intermediate) outputs_coord = self.conditional_detr.bbox_predictor(intermediate).sigmoid() loss, loss_dict, auxiliary_outputs = self.loss_function( logits, labels, device, pred_boxes, pred_masks, self.config, outputs_class, outputs_coord ) return ConditionalDetrSegmentationOutput( loss=loss, loss_dict=loss_dict, logits=logits, pred_boxes=pred_boxes, pred_masks=pred_masks, auxiliary_outputs=auxiliary_outputs, last_hidden_state=decoder_outputs.last_hidden_state, decoder_hidden_states=decoder_outputs.hidden_states, decoder_attentions=decoder_outputs.attentions, cross_attentions=decoder_outputs.cross_attentions, encoder_last_hidden_state=encoder_outputs.last_hidden_state, encoder_hidden_states=encoder_outputs.hidden_states, encoder_attentions=encoder_outputs.attentions, ) __all__ = [ "ConditionalDetrForObjectDetection", "ConditionalDetrForSegmentation", "ConditionalDetrModel", "ConditionalDetrPreTrainedModel", ]