# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/deimv2/modular_deimv2.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_deimv2.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # Copyright 2026 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 import torch.nn as nn import torch.nn.functional as F from torch import Tensor from ... import initialization as init from ...activations import ACT2CLS from ...backbone_utils import load_backbone from ...image_transforms import center_to_corners_format, corners_to_center_format from ...integrations import use_kernel_forward_from_hub from ...modeling_outputs import ModelOutput 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 TransformersKwargs, auto_docstring, torch_compilable_check, torch_int from ...utils.generic import can_return_tuple, merge_with_config_defaults from ...utils.output_capturing import OutputRecorder, capture_outputs from .configuration_deimv2 import Deimv2Config @dataclass @auto_docstring( custom_intro=""" Base class for outputs of the Deimv2Decoder. This class adds two attributes to BaseModelOutputWithCrossAttentions, namely: - a stacked tensor of intermediate decoder hidden states (i.e. the output of each decoder layer) - a stacked tensor of intermediate reference points. """ ) class Deimv2DecoderOutput(ModelOutput): r""" intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`): Stacked intermediate hidden states (output of each layer of the decoder). intermediate_logits (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, sequence_length, config.num_labels)`): Stacked intermediate logits (logits of each layer of the decoder). intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, sequence_length, hidden_size)`): Stacked intermediate reference points (reference points of each layer of the decoder). intermediate_predicted_corners (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): Stacked intermediate predicted corners (predicted corners of each layer of the decoder). initial_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): Stacked initial reference points (initial reference points of each layer of the decoder). 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. """ last_hidden_state: torch.FloatTensor | None = None intermediate_hidden_states: torch.FloatTensor | None = None intermediate_logits: torch.FloatTensor | None = None intermediate_reference_points: torch.FloatTensor | None = None intermediate_predicted_corners: torch.FloatTensor | None = None initial_reference_points: torch.FloatTensor | None = None hidden_states: tuple[torch.FloatTensor] | None = None attentions: tuple[torch.FloatTensor] | None = None cross_attentions: tuple[torch.FloatTensor] | None = None @dataclass @auto_docstring( custom_intro=""" Base class for outputs of the RT-DETR encoder-decoder model. """ ) class Deimv2ModelOutput(ModelOutput): r""" last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_queries, 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 `(batch_size, config.decoder_layers, num_queries, hidden_size)`): Stacked intermediate hidden states (output of each layer of the decoder). intermediate_logits (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, sequence_length, config.num_labels)`): Stacked intermediate logits (logits of each layer of the decoder). intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): Stacked intermediate reference points (reference points of each layer of the decoder). intermediate_predicted_corners (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): Stacked intermediate predicted corners (predicted corners of each layer of the decoder). initial_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`): Initial reference points used for the first decoder layer. init_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`): Initial reference points sent through the Transformer decoder. enc_topk_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`): Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are picked as region proposals in the encoder stage. Output of bounding box binary classification (i.e. foreground and background). enc_topk_bboxes (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`): Logits of predicted bounding boxes coordinates in the encoder stage. enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are picked as region proposals in the first stage. Output of bounding box binary classification (i.e. foreground and background). enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): Logits of predicted bounding boxes coordinates in the first stage. denoising_meta_values (`dict`): Extra dictionary for the denoising related values. """ last_hidden_state: torch.FloatTensor | None = None intermediate_hidden_states: torch.FloatTensor | None = None intermediate_logits: torch.FloatTensor | None = None intermediate_reference_points: torch.FloatTensor | None = None intermediate_predicted_corners: torch.FloatTensor | None = None initial_reference_points: 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 init_reference_points: torch.FloatTensor | None = None enc_topk_logits: torch.FloatTensor | None = None enc_topk_bboxes: torch.FloatTensor | None = None enc_outputs_class: torch.FloatTensor | None = None enc_outputs_coord_logits: torch.FloatTensor | None = None denoising_meta_values: dict | None = None @dataclass @auto_docstring( custom_intro=""" Output type for DEIMv2 encoder modules (HybridEncoder and LiteEncoder). Attentions are only available for HybridEncoder variants with AIFI layers. """ ) class Deimv2EncoderOutput(ModelOutput): r""" feature_maps (`list[torch.FloatTensor]`): List of multi-scale feature maps from the encoder, one per feature level. """ feature_maps: list[torch.FloatTensor] = None hidden_states: tuple[torch.FloatTensor, ...] | None = None attentions: tuple[torch.FloatTensor, ...] | None = None @use_kernel_forward_from_hub("RMSNorm") class Deimv2RMSNorm(nn.Module): def __init__(self, hidden_size, eps: float = 1e-6) -> None: """ Deimv2RMSNorm is equivalent to T5LayerNorm """ super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) return self.weight * hidden_states.to(input_dtype) def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" class Deimv2SwiGLUFFN(nn.Module): def __init__(self, config: Deimv2Config): super().__init__() hidden_features = config.decoder_ffn_dim // 2 self.gate_proj = nn.Linear(config.d_model, hidden_features, bias=True) self.up_proj = nn.Linear(config.d_model, hidden_features, bias=True) self.down_proj = nn.Linear(hidden_features, config.d_model, bias=True) self.act_fn = nn.SiLU() def forward(self, x): down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) return down_proj class Deimv2Gate(nn.Module): def __init__(self, d_model: int): super().__init__() self.gate = nn.Linear(2 * d_model, 2 * d_model) self.norm = Deimv2RMSNorm(d_model) def forward(self, second_residual: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: gate_input = torch.cat([second_residual, hidden_states], dim=-1) gates = torch.sigmoid(self.gate(gate_input)) gate1, gate2 = gates.chunk(2, dim=-1) hidden_states = self.norm(gate1 * second_residual + gate2 * hidden_states) return hidden_states class Deimv2MLP(nn.Module): def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int, act: str = "relu"): super().__init__() self.num_layers = num_layers hidden_dims = [hidden_dim] * (num_layers - 1) input_dims = [input_dim] + hidden_dims output_dims = hidden_dims + [output_dim] self.layers = nn.ModuleList(nn.Linear(in_dim, out_dim) for in_dim, out_dim in zip(input_dims, output_dims)) self.act = ACT2CLS[act]() def forward(self, stat_features: torch.Tensor) -> torch.Tensor: for i, layer in enumerate(self.layers): stat_features = self.act(layer(stat_features)) if i < self.num_layers - 1 else layer(stat_features) return stat_features def multi_scale_deformable_attention_v2( value: Tensor, value_spatial_shapes: Tensor, sampling_locations: Tensor, attention_weights: Tensor, num_points_list: list[int], method="default", ) -> Tensor: batch_size, _, num_heads, hidden_dim = value.shape _, num_queries, num_heads, num_levels, num_points = sampling_locations.shape value_list = ( value.permute(0, 2, 3, 1) .flatten(0, 1) .split([height * width for height, width in value_spatial_shapes], dim=-1) ) # sampling_offsets [8, 480, 8, 12, 2] if method == "default": sampling_grids = 2 * sampling_locations - 1 elif method == "discrete": sampling_grids = sampling_locations sampling_grids = sampling_grids.permute(0, 2, 1, 3, 4).flatten(0, 1) sampling_grids = sampling_grids.split(num_points_list, dim=-2) sampling_value_list = [] for level_id, (height, width) in enumerate(value_spatial_shapes): # batch_size, height*width, num_heads, hidden_dim # -> batch_size, height*width, num_heads*hidden_dim # -> batch_size, num_heads*hidden_dim, height*width # -> batch_size*num_heads, hidden_dim, height, width value_l_ = value_list[level_id].reshape(batch_size * num_heads, hidden_dim, height, width) # batch_size, num_queries, num_heads, num_points, 2 # -> batch_size, num_heads, num_queries, num_points, 2 # -> batch_size*num_heads, num_queries, num_points, 2 sampling_grid_l_ = sampling_grids[level_id] # batch_size*num_heads, hidden_dim, num_queries, num_points if method == "default": sampling_value_l_ = nn.functional.grid_sample( value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False ) elif method == "discrete": sampling_coord = (sampling_grid_l_ * torch.tensor([[width, height]], device=value.device) + 0.5).to( torch.int64 ) # Separate clamping for x and y coordinates sampling_coord_x = sampling_coord[..., 0].clamp(0, width - 1) sampling_coord_y = sampling_coord[..., 1].clamp(0, height - 1) # Combine the clamped coordinates sampling_coord = torch.stack([sampling_coord_x, sampling_coord_y], dim=-1) sampling_coord = sampling_coord.reshape(batch_size * num_heads, num_queries * num_points_list[level_id], 2) sampling_idx = ( torch.arange(sampling_coord.shape[0], device=value.device) .unsqueeze(-1) .repeat(1, sampling_coord.shape[1]) ) sampling_value_l_ = value_l_[sampling_idx, :, sampling_coord[..., 1], sampling_coord[..., 0]] sampling_value_l_ = sampling_value_l_.permute(0, 2, 1).reshape( batch_size * num_heads, hidden_dim, num_queries, num_points_list[level_id] ) sampling_value_list.append(sampling_value_l_) # (batch_size, num_queries, num_heads, num_levels, num_points) # -> (batch_size, num_heads, num_queries, num_levels, num_points) # -> (batch_size, num_heads, 1, num_queries, num_levels*num_points) attention_weights = attention_weights.permute(0, 2, 1, 3).reshape( batch_size * num_heads, 1, num_queries, sum(num_points_list) ) output = ( (torch.concat(sampling_value_list, dim=-1) * attention_weights) .sum(-1) .view(batch_size, num_heads * hidden_dim, num_queries) ) return output.transpose(1, 2).contiguous() class Deimv2MultiscaleDeformableAttention(nn.Module): def __init__(self, config: Deimv2Config): """ D-Fine version of multiscale deformable attention """ super().__init__() self.d_model = config.d_model self.n_heads = config.decoder_attention_heads self.n_levels = config.num_feature_levels self.offset_scale = config.decoder_offset_scale self.decoder_method = config.decoder_method self.n_points = config.decoder_n_points if isinstance(self.n_points, list): num_points_list = self.n_points else: num_points_list = [self.n_points for _ in range(self.n_levels)] self.num_points_list = num_points_list num_points_scale = [1 / n for n in self.num_points_list for _ in range(n)] self.register_buffer("num_points_scale", torch.tensor(num_points_scale, dtype=torch.float32)) self.total_points = self.n_heads * sum(self.num_points_list) self.sampling_offsets = nn.Linear(self.d_model, self.total_points * 2) self.attention_weights = nn.Linear(self.d_model, self.total_points) self.ms_deformable_attn_core = multi_scale_deformable_attention_v2 def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, reference_points=None, encoder_hidden_states=None, spatial_shapes=None, spatial_shapes_list=None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor, torch.Tensor]: batch_size, num_queries, _ = hidden_states.shape batch_size, sequence_length, _ = encoder_hidden_states.shape torch_compilable_check( (spatial_shapes[:, 0] * spatial_shapes[:, 1]).sum() == sequence_length, "Make sure to align the spatial shapes with the sequence length of the encoder hidden states", ) # Reshape for multi-head attention value = encoder_hidden_states.reshape(batch_size, sequence_length, self.n_heads, self.d_model // self.n_heads) if attention_mask is not None: value = value.masked_fill(~attention_mask[..., None], float(0)) sampling_offsets: torch.Tensor = self.sampling_offsets(hidden_states) sampling_offsets = sampling_offsets.reshape( batch_size, num_queries, self.n_heads, sum(self.num_points_list), 2 ) attention_weights = self.attention_weights(hidden_states).reshape( batch_size, num_queries, self.n_heads, sum(self.num_points_list) ) attention_weights = F.softmax(attention_weights, dim=-1) if reference_points.shape[-1] == 2: offset_normalizer = torch.tensor(spatial_shapes) offset_normalizer = offset_normalizer.flip([1]).reshape(1, 1, 1, self.n_levels, 1, 2) sampling_locations = ( reference_points.reshape(batch_size, sequence_length, 1, self.n_levels, 1, 2) + sampling_offsets / offset_normalizer ) elif reference_points.shape[-1] == 4: # reference_points [8, 480, None, 1, 4] # sampling_offsets [8, 480, 8, 12, 2] num_points_scale = self.num_points_scale.to(dtype=hidden_states.dtype).unsqueeze(-1) offset = sampling_offsets * num_points_scale * reference_points[:, :, None, :, 2:] * self.offset_scale sampling_locations = reference_points[:, :, None, :, :2] + offset else: raise ValueError( f"Last dim of reference_points must be 2 or 4, but get {reference_points.shape[-1]} instead." ) output = self.ms_deformable_attn_core( value, spatial_shapes_list, sampling_locations, attention_weights, self.num_points_list, self.decoder_method, ) return output, attention_weights class Deimv2ConvNormLayer(nn.Module): def __init__( self, config: Deimv2Config, in_channels: int, out_channels: int, kernel_size: int, stride: int, groups: int = 1, padding: int | None = None, activation: str | None = None, ): super().__init__() self.conv = nn.Conv2d( in_channels, out_channels, kernel_size, stride, groups=groups, padding=(kernel_size - 1) // 2 if padding is None else padding, bias=False, ) self.norm = nn.BatchNorm2d(out_channels, config.batch_norm_eps) self.activation = nn.Identity() if activation is None else ACT2CLS[activation]() def forward(self, hidden_state): hidden_state = self.conv(hidden_state) hidden_state = self.norm(hidden_state) hidden_state = self.activation(hidden_state) return hidden_state class Deimv2RepVggBlock(nn.Module): """ RepVGG architecture block introduced by the work "RepVGG: Making VGG-style ConvNets Great Again". """ def __init__(self, config: Deimv2Config, in_channels: int, out_channels: int): super().__init__() activation = config.activation_function hidden_channels = in_channels self.conv1 = Deimv2ConvNormLayer(config, hidden_channels, out_channels, 3, 1, padding=1) self.conv2 = Deimv2ConvNormLayer(config, hidden_channels, out_channels, 1, 1, padding=0) self.activation = nn.Identity() if activation is None else ACT2CLS[activation]() def forward(self, x): y = self.conv1(x) + self.conv2(x) return self.activation(y) class Deimv2CSPRepLayer(nn.Module): """ Cross Stage Partial (CSP) network layer with RepVGG blocks. Differs from DFineCSPRepLayer: uses a single conv that splits into residual + processing path (instead of two separate convs), and has an optional trailing conv controlled by `encoder_has_trailing_conv`. """ def __init__( self, config: Deimv2Config, in_channels: int, out_channels: int, num_blocks: int, expansion: float = 1.0 ): super().__init__() activation = config.activation_function hidden_channels = int(out_channels * expansion) self.conv1 = Deimv2ConvNormLayer(config, in_channels, hidden_channels * 2, 1, 1, activation=activation) self.bottlenecks = nn.ModuleList( [Deimv2RepVggBlock(config, hidden_channels, hidden_channels) for _ in range(num_blocks)] ) self.conv2 = ( Deimv2ConvNormLayer(config, hidden_channels, out_channels, 3, 1, activation=activation) if config.encoder_has_trailing_conv else nn.Identity() ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: residual, hidden_states = self.conv1(hidden_states).chunk(2, dim=1) for bottleneck in self.bottlenecks: hidden_states = bottleneck(hidden_states) return self.conv2(residual + hidden_states) class Deimv2RepNCSPELAN5(nn.Module): """ Rep(VGG) N(etwork) CSP (Cross Stage Partial) ELAN (Efficient Layer Aggregation Network) block. Similar to DFineRepNCSPELAN4 but without intermediate convolutions between CSP branches, resulting in a simpler 4-way concatenation (2 split halves + 2 CSP branches) instead of D-FINE's 4-branch design with interleaved convolutions. """ def __init__(self, config: Deimv2Config, numb_blocks: int = 3): super().__init__() activation = config.activation_function in_channels = config.encoder_hidden_dim out_channels = config.encoder_hidden_dim split_channels = config.encoder_hidden_dim * 2 csp_channels = round(config.hidden_expansion * config.encoder_hidden_dim // 2) self.conv1 = Deimv2ConvNormLayer(config, in_channels, split_channels, 1, 1, activation=activation) self.csp_rep1 = Deimv2CSPRepLayer(config, split_channels // 2, csp_channels, num_blocks=numb_blocks) self.csp_rep2 = Deimv2CSPRepLayer(config, csp_channels, csp_channels, num_blocks=numb_blocks) self.conv2 = Deimv2ConvNormLayer( config, split_channels + (2 * csp_channels), out_channels, 1, 1, activation=activation ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states_1, hidden_states_2 = self.conv1(hidden_states).chunk(2, dim=1) hidden_states_3 = self.csp_rep1(hidden_states_2) hidden_states_4 = self.csp_rep2(hidden_states_3) merged_hidden_states = torch.cat([hidden_states_1, hidden_states_2, hidden_states_3, hidden_states_4], dim=1) return self.conv2(merged_hidden_states) class Deimv2SCDown(nn.Module): def __init__(self, config: Deimv2Config, kernel_size: int, stride: int): super().__init__() self.conv1 = Deimv2ConvNormLayer(config, config.encoder_hidden_dim, config.encoder_hidden_dim, 1, 1) self.conv2 = Deimv2ConvNormLayer( config, config.encoder_hidden_dim, config.encoder_hidden_dim, kernel_size, stride, config.encoder_hidden_dim, ) def forward(self, input_features: torch.Tensor) -> torch.Tensor: input_features = self.conv1(input_features) input_features = self.conv2(input_features) return input_features 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 Deimv2SelfAttention(nn.Module): """ Multi-headed self-attention from 'Attention Is All You Need' paper. In DEIMV2, position embeddings are added to both queries and keys (but not values) in self-attention. """ def __init__( self, config: Deimv2Config, 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 Deimv2EncoderLayer(nn.Module): def __init__(self, config: Deimv2Config): super().__init__() self.normalize_before = config.normalize_before self.hidden_size = config.encoder_hidden_dim # self-attention self.self_attn = Deimv2SelfAttention( config=config, hidden_size=self.hidden_size, num_attention_heads=config.num_attention_heads, dropout=config.dropout, ) self.self_attn_layer_norm = nn.LayerNorm(self.hidden_size, eps=config.layer_norm_eps) self.dropout = config.dropout self.mlp = Deimv2MLP( self.hidden_size, config.encoder_ffn_dim, self.hidden_size, 2, config.encoder_activation_function ) self.final_layer_norm = nn.LayerNorm(self.hidden_size, eps=config.layer_norm_eps) 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 if self.normalize_before: hidden_states = self.self_attn_layer_norm(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 if not self.normalize_before: hidden_states = self.self_attn_layer_norm(hidden_states) if self.normalize_before: hidden_states = self.final_layer_norm(hidden_states) residual = hidden_states hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states if not self.normalize_before: 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 Deimv2SinePositionEmbedding(nn.Module): """ 2D sinusoidal position embedding used in RT-DETR hybrid encoder. """ def __init__(self, embed_dim: int = 256, temperature: int = 10000): super().__init__() self.embed_dim = embed_dim self.temperature = temperature @compile_compatible_method_lru_cache(maxsize=32) def forward( self, width: int, height: int, device: torch.device | str, dtype: torch.dtype, ) -> torch.Tensor: """ Generate 2D sinusoidal position embeddings. Returns: Position embeddings of shape (1, height*width, embed_dim) """ grid_w = torch.arange(torch_int(width), device=device).to(dtype) grid_h = torch.arange(torch_int(height), device=device).to(dtype) grid_w, grid_h = torch.meshgrid(grid_w, grid_h, indexing="xy") if self.embed_dim % 4 != 0: raise ValueError("Embed dimension must be divisible by 4 for 2D sin-cos position embedding") pos_dim = self.embed_dim // 4 omega = torch.arange(pos_dim, device=device).to(dtype) / pos_dim omega = 1.0 / (self.temperature**omega) out_w = grid_w.flatten()[..., None] @ omega[None] out_h = grid_h.flatten()[..., None] @ omega[None] return torch.concat([out_h.sin(), out_h.cos(), out_w.sin(), out_w.cos()], dim=1)[None, :, :] class Deimv2AIFILayer(nn.Module): """ AIFI (Attention-based Intra-scale Feature Interaction) layer used in RT-DETR hybrid encoder. """ def __init__(self, config: Deimv2Config): super().__init__() self.config = config self.encoder_hidden_dim = config.encoder_hidden_dim self.eval_size = config.eval_size self.position_embedding = Deimv2SinePositionEmbedding( embed_dim=self.encoder_hidden_dim, temperature=config.positional_encoding_temperature, ) self.layers = nn.ModuleList([Deimv2EncoderLayer(config) for _ in range(config.encoder_layers)]) def forward( self, hidden_states: torch.Tensor, **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: """ Args: hidden_states (`torch.FloatTensor` of shape `(batch_size, channels, height, width)`): Feature map to process. """ batch_size = hidden_states.shape[0] height, width = hidden_states.shape[2:] hidden_states = hidden_states.flatten(2).permute(0, 2, 1) if self.training or self.eval_size is None: pos_embed = self.position_embedding( width=width, height=height, device=hidden_states.device, dtype=hidden_states.dtype, ) else: pos_embed = None for layer in self.layers: hidden_states = layer( hidden_states, attention_mask=None, spatial_position_embeddings=pos_embed, **kwargs, ) hidden_states = ( hidden_states.permute(0, 2, 1).reshape(batch_size, self.encoder_hidden_dim, height, width).contiguous() ) return hidden_states class Deimv2SpatialTuningAdapter(nn.Module): def __init__(self, config: Deimv2Config): super().__init__() inplanes = config.spatial_tuning_adapter_inplanes self.stem_conv = Deimv2ConvNormLayer(config, 3, inplanes, 3, 2, activation="gelu") self.stem_pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) self.conv2 = Deimv2ConvNormLayer(config, inplanes, 2 * inplanes, 3, 2) self.conv3 = Deimv2ConvNormLayer(config, 2 * inplanes, 4 * inplanes, 3, 2) self.conv4 = Deimv2ConvNormLayer(config, 4 * inplanes, 4 * inplanes, 3, 2) self.act_fn = nn.GELU() def forward(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: hidden_states_1 = self.stem_pool(self.stem_conv(pixel_values)) hidden_states_2 = self.conv2(hidden_states_1) hidden_states_3 = self.conv3(self.act_fn(hidden_states_2)) hidden_states_4 = self.conv4(self.act_fn(hidden_states_3)) return hidden_states_2, hidden_states_3, hidden_states_4 class Deimv2FrozenBatchNorm2d(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 `Deimv2FrozenBatchNorm2d`. Args: model (torch.nn.Module): input model """ for name, module in model.named_children(): if isinstance(module, nn.BatchNorm2d): new_module = Deimv2FrozenBatchNorm2d(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 Deimv2ConvEncoder(nn.Module): """ Convolutional backbone using the modeling_deimv2_resnet.py. nn.BatchNorm2d layers are replaced by Deimv2FrozenBatchNorm2d as defined above. https://github.com/lyuwenyu/RT-DETR/blob/main/Deimv2_pytorch/src/nn/backbone/presnet.py#L142 """ def __init__(self, config): super().__init__() backbone = load_backbone(config) if config.freeze_backbone_batch_norms: # replace batch norm by frozen batch norm with torch.no_grad(): replace_batch_norm(backbone) self.model = backbone self.intermediate_channel_sizes = self.model.channels self.encoder_input_proj = nn.ModuleList( [ Deimv2ConvNormLayer(config, in_channel, config.encoder_hidden_dim, 1, 1) if config.encoder_type != "lite" else nn.Identity() for in_channel in self.intermediate_channel_sizes ] ) def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> list[torch.Tensor]: features = self.model(pixel_values, **kwargs).feature_maps return [proj(feat) for proj, feat in zip(self.encoder_input_proj, features)] class Deimv2DINOv3ConvEncoder(nn.Module): def __init__(self, config: Deimv2Config): super().__init__() self.backbone = load_backbone(config) self.spatial_tuning_adapter = Deimv2SpatialTuningAdapter(config) embed_dim = config.backbone_config.hidden_size hidden_dim = config.encoder_hidden_dim spatial_tuning_adapter_channels = config.spatial_tuning_adapter_inplanes self.fusion_proj = nn.ModuleList( [ Deimv2ConvNormLayer(config, embed_dim + spatial_tuning_adapter_channels * 2, hidden_dim, 1, 1), Deimv2ConvNormLayer(config, embed_dim + spatial_tuning_adapter_channels * 4, hidden_dim, 1, 1), Deimv2ConvNormLayer(config, embed_dim + spatial_tuning_adapter_channels * 4, hidden_dim, 1, 1), ] ) def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> list[torch.Tensor]: backbone_output = self.backbone(pixel_values, **kwargs) feature_maps = backbone_output.feature_maps patch_size = self.backbone.config.patch_size height_patches = pixel_values.shape[2] // patch_size width_patches = pixel_values.shape[3] // patch_size semantic_features = [] num_scales = len(feature_maps) for i, feat in enumerate(feature_maps): resize_height = int(height_patches * 2 ** (num_scales - 2 - i)) resize_width = int(width_patches * 2 ** (num_scales - 2 - i)) spatial = F.interpolate(feat, size=[resize_height, resize_width], mode="bilinear", align_corners=False) semantic_features.append(spatial) detail_features = self.spatial_tuning_adapter(pixel_values) outputs = [] for i, (semantic_feature, detail_feature) in enumerate(zip(semantic_features, detail_features)): fused = torch.cat([semantic_feature, detail_feature], dim=1) outputs.append(self.fusion_proj[i](fused)) return outputs class Deimv2Integral(nn.Module): """ A static layer that calculates integral results from a distribution. This layer computes the target location using the formula: `sum{Pr(n) * W(n)}`, where Pr(n) is the softmax probability vector representing the discrete distribution, and W(n) is the non-uniform Weighting Function. Args: max_num_bins (int): Max number of the discrete bins. Default is 32. It can be adjusted based on the dataset or task requirements. """ def __init__(self, config: Deimv2Config): super().__init__() self.max_num_bins = config.max_num_bins def forward(self, pred_corners: torch.Tensor, project: torch.Tensor) -> torch.Tensor: batch_size, num_queries, _ = pred_corners.shape pred_corners = F.softmax(pred_corners.reshape(-1, self.max_num_bins + 1), dim=1) pred_corners = F.linear(pred_corners, project.to(pred_corners.device)).reshape(-1, 4) pred_corners = pred_corners.reshape(batch_size, num_queries, -1) return pred_corners class Deimv2LQE(nn.Module): def __init__(self, config: Deimv2Config): super().__init__() self.top_prob_values = config.top_prob_values self.max_num_bins = config.max_num_bins self.reg_conf = Deimv2MLP(4 * (self.top_prob_values + 1), config.lqe_hidden_dim, 1, config.lqe_layers) def forward(self, scores: torch.Tensor, pred_corners: torch.Tensor) -> torch.Tensor: batch_size, length, _ = pred_corners.size() prob = F.softmax(pred_corners.reshape(batch_size, length, 4, self.max_num_bins + 1), dim=-1) prob_topk, _ = prob.topk(self.top_prob_values, dim=-1) stat = torch.cat([prob_topk, prob_topk.mean(dim=-1, keepdim=True)], dim=-1) quality_score = self.reg_conf(stat.reshape(batch_size, length, -1)) scores = scores + quality_score return scores class Deimv2DecoderLayer(nn.Module): def __init__(self, config: Deimv2Config): super().__init__() self.hidden_size = config.d_model # self-attention self.self_attn = Deimv2SelfAttention( 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 = Deimv2RMSNorm(config.d_model) self.encoder_attn = Deimv2MultiscaleDeformableAttention(config=config) self.mlp = Deimv2SwiGLUFFN(config) self.final_layer_norm = Deimv2RMSNorm(config.d_model) self.gateway = Deimv2Gate(config.d_model) if config.use_gateway else None self.use_gateway = config.use_gateway self.encoder_attn_layer_norm = None if config.use_gateway else Deimv2RMSNorm(config.d_model) def forward( self, hidden_states: torch.Tensor, position_embeddings: torch.Tensor | None = None, reference_points: torch.Tensor | None = None, spatial_shapes: torch.Tensor | None = None, spatial_shapes_list: list[tuple[int, int]] | None = None, encoder_hidden_states: torch.Tensor | None = None, encoder_attention_mask: 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)`. object_queries_position_embeddings (`torch.FloatTensor`, *optional*): Position embeddings for the object query slots. These are added to both queries and keys in the self-attention layer (not values). reference_points (`torch.FloatTensor`, *optional*): Reference points. spatial_shapes (`torch.LongTensor`, *optional*): Spatial shapes. level_start_index (`torch.LongTensor`, *optional*): Level start index. encoder_hidden_states (`torch.FloatTensor`): cross attention input to the layer of shape `(batch, seq_len, hidden_size)` 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. """ residual = hidden_states # Self Attention hidden_states, _ = self.self_attn( hidden_states=hidden_states, attention_mask=encoder_attention_mask, position_embeddings=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 # Cross-Attention hidden_states = hidden_states if position_embeddings is None else hidden_states + position_embeddings hidden_states, _ = self.encoder_attn( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, reference_points=reference_points, spatial_shapes=spatial_shapes, spatial_shapes_list=spatial_shapes_list, ) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) if self.gateway is not None: hidden_states = self.gateway(residual, hidden_states) else: 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 @auto_docstring class Deimv2PreTrainedModel(PreTrainedModel): config: Deimv2Config base_model_prefix = "deimv2" main_input_name = "pixel_values" input_modalities = ("image",) _no_split_modules = [r"Deimv2HybridEncoder", r"Deimv2LiteEncoder", r"Deimv2DecoderLayer"] _supports_sdpa = True _supports_flash_attn = True _supports_attention_backend = True _supports_flex_attn = True @torch.no_grad() def _init_weights(self, module): """Initialize the weights""" super()._init_weights(module) # initialize linear layer bias value according to a given probability value. if isinstance(module, (Deimv2ForObjectDetection, Deimv2Decoder)): if module.class_embed is not None: for layer in module.class_embed: prior_prob = self.config.initializer_bias_prior_prob or 1 / (self.config.num_labels + 1) bias = float(-math.log((1 - prior_prob) / prior_prob)) init.xavier_uniform_(layer.weight) init.constant_(layer.bias, bias) if module.bbox_embed is not None: for layer in module.bbox_embed: init.constant_(layer.layers[-1].weight, 0) init.constant_(layer.layers[-1].bias, 0) if hasattr(module, "reg_scale"): init.constant_(module.reg_scale, self.config.reg_scale) if hasattr(module, "up"): init.constant_(module.up, self.config.up) if isinstance(module, Deimv2MultiscaleDeformableAttention): init.constant_(module.sampling_offsets.weight, 0.0) default_dtype = torch.get_default_dtype() thetas = torch.arange(module.n_heads, dtype=torch.int64).to(default_dtype) * ( 2.0 * math.pi / module.n_heads ) grid_init = torch.stack([thetas.cos(), thetas.sin()], -1) grid_init = grid_init / grid_init.abs().max(-1, keepdim=True).values grid_init = grid_init.reshape(module.n_heads, 1, 2).tile([1, sum(module.num_points_list), 1]) scaling = torch.concat([torch.arange(1, n + 1) for n in module.num_points_list]).reshape(1, -1, 1) grid_init *= scaling init.copy_(module.sampling_offsets.bias, grid_init.flatten()) init.constant_(module.attention_weights.weight, 0.0) init.constant_(module.attention_weights.bias, 0.0) num_points_scale = [1 / n for n in module.num_points_list for _ in range(n)] init.copy_(module.num_points_scale, torch.tensor(num_points_scale, dtype=torch.float32)) if isinstance(module, Deimv2Model): prior_prob = self.config.initializer_bias_prior_prob or 1 / (self.config.num_labels + 1) bias = float(-math.log((1 - prior_prob) / prior_prob)) init.xavier_uniform_(module.enc_score_head.weight) init.constant_(module.enc_score_head.bias, bias) if isinstance(module, Deimv2Gate): bias = float(-math.log((1 - 0.5) / 0.5)) init.constant_(module.gate.bias, bias) init.constant_(module.gate.weight, 0) if isinstance(module, Deimv2LQE): init.constant_(module.reg_conf.layers[-1].bias, 0) init.constant_(module.reg_conf.layers[-1].weight, 0) if hasattr(module, "weight_embedding") and self.config.learn_initial_query: init.xavier_uniform_(module.weight_embedding.weight) if hasattr(module, "denoising_class_embed") and self.config.num_denoising > 0: init.xavier_uniform_(module.denoising_class_embed.weight) if isinstance(module, Deimv2SwiGLUFFN): init.xavier_uniform_(module.gate_proj.weight) init.constant_(module.gate_proj.bias, 0) init.xavier_uniform_(module.up_proj.weight) init.constant_(module.up_proj.bias, 0) init.xavier_uniform_(module.down_proj.weight) init.constant_(module.down_proj.bias, 0) class Deimv2LiteEncoder(Deimv2PreTrainedModel): # LiteEncoder has no transformer layers, so hidden_states are recorded from the conv projections. _can_record_outputs = { "hidden_states": [ OutputRecorder(Deimv2ConvNormLayer, layer_name="input_proj"), OutputRecorder(Deimv2ConvNormLayer, layer_name="bi_fusion_conv"), ], } def __init__(self, config: Deimv2Config): super().__init__(config) hidden_dim = config.encoder_hidden_dim activation = config.activation_function self.input_proj = nn.ModuleList( [Deimv2ConvNormLayer(config, in_channel, hidden_dim, 1, 1) for in_channel in config.encoder_in_channels] ) self.down_pool1 = nn.AvgPool2d(kernel_size=3, stride=2, padding=1) self.down_conv1 = Deimv2ConvNormLayer(config, hidden_dim, hidden_dim, 1, 1, activation=activation) self.down_pool2 = nn.AvgPool2d(kernel_size=3, stride=2, padding=1) self.down_conv2 = Deimv2ConvNormLayer(config, hidden_dim, hidden_dim, 1, 1, activation=activation) self.bi_fusion_conv = Deimv2ConvNormLayer(config, hidden_dim, hidden_dim, 1, 1, activation=activation) num_blocks = round(3 * config.depth_mult) self.fpn_block = Deimv2RepNCSPELAN5(config, numb_blocks=num_blocks) self.pan_block = Deimv2RepNCSPELAN5(config, numb_blocks=num_blocks) self.post_init() @merge_with_config_defaults @capture_outputs def forward(self, inputs_embeds: list[torch.Tensor], **kwargs: Unpack[TransformersKwargs]) -> Deimv2EncoderOutput: projected_features = [self.input_proj[i](feature) for i, feature in enumerate(inputs_embeds)] projected_features.append(self.down_conv1(self.down_pool1(projected_features[-1]))) projected_features[-1] = self.bi_fusion_conv( projected_features[-1] + F.adaptive_avg_pool2d(projected_features[-1], 1) ) outputs = [] fused_feature = projected_features[0] + F.interpolate(projected_features[1], scale_factor=2.0, mode="nearest") outputs.append(self.fpn_block(fused_feature)) fused_feature = projected_features[1] + self.down_conv2(self.down_pool2(outputs[-1])) outputs.append(self.pan_block(fused_feature)) return Deimv2EncoderOutput(feature_maps=outputs) def fuse_feature_maps(feature_map_1: torch.Tensor, feature_map_2: torch.Tensor, fuse_op: str = "sum") -> torch.Tensor: """Fuses two feature maps via element-wise sum or channel-wise concatenation.""" if fuse_op == "sum": return feature_map_1 + feature_map_2 return torch.cat([feature_map_1, feature_map_2], dim=1) class Deimv2HybridEncoder(Deimv2PreTrainedModel): """ DEIMv2 variant of DFineHybridEncoder. Uses element-wise sum fusion (`fuse_feature_maps`) instead of D-FINE's channel concatenation, Deimv2RepNCSPELAN5 (simplified 4-way concat) instead of DFineRepNCSPELAN4, and returns Deimv2EncoderOutput with feature_maps instead of BaseModelOutput with last_hidden_state. """ _can_record_outputs = { "hidden_states": Deimv2AIFILayer, "attentions": Deimv2SelfAttention, } def __init__(self, config: Deimv2Config): super().__init__(config) self.config = config self.in_channels = config.encoder_in_channels self.num_fpn_stages = len(self.in_channels) - 1 self.feat_strides = config.feat_strides self.encoder_hidden_dim = config.encoder_hidden_dim self.encode_proj_layers = config.encode_proj_layers self.positional_encoding_temperature = config.positional_encoding_temperature self.eval_size = config.eval_size self.out_channels = [self.encoder_hidden_dim for _ in self.in_channels] self.out_strides = self.feat_strides self.fuse_op = config.encoder_fuse_op self.aifi = nn.ModuleList([Deimv2AIFILayer(config) for _ in range(len(self.encode_proj_layers))]) self.lateral_convs = nn.ModuleList() self.fpn_blocks = nn.ModuleList() for _ in range(len(self.in_channels) - 1, 0, -1): self.lateral_convs.append( Deimv2ConvNormLayer(config, self.encoder_hidden_dim, self.encoder_hidden_dim, 1, 1) ) num_blocks = round(3 * config.depth_mult) self.fpn_blocks.append(Deimv2RepNCSPELAN5(config, numb_blocks=num_blocks)) self.downsample_convs = nn.ModuleList() self.pan_blocks = nn.ModuleList() for _ in range(len(self.in_channels) - 1): self.downsample_convs.append(Deimv2SCDown(config, 3, 2)) num_blocks = round(3 * config.depth_mult) self.pan_blocks.append(Deimv2RepNCSPELAN5(config, numb_blocks=num_blocks)) self.post_init() @merge_with_config_defaults @capture_outputs(tie_last_hidden_states=False) def forward( self, inputs_embeds: list[torch.Tensor] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> Deimv2EncoderOutput: r""" Args: inputs_embeds (`list[torch.FloatTensor]`): Multi-scale feature maps from the backbone (one tensor per feature level) passed to the encoder. """ feature_maps = inputs_embeds if self.config.encoder_layers > 0: for i, enc_ind in enumerate(self.encode_proj_layers): feature_maps[enc_ind] = self.aifi[i](feature_maps[enc_ind], **kwargs) # top-down FPN fpn_feature_maps = [feature_maps[-1]] for idx, (lateral_conv, fpn_block) in enumerate(zip(self.lateral_convs, self.fpn_blocks)): backbone_feature_map = feature_maps[self.num_fpn_stages - idx - 1] top_fpn_feature_map = fpn_feature_maps[-1] top_fpn_feature_map = lateral_conv(top_fpn_feature_map) fpn_feature_maps[-1] = top_fpn_feature_map top_fpn_feature_map = F.interpolate(top_fpn_feature_map, scale_factor=2.0, mode="nearest") fused_feature_map = fuse_feature_maps(top_fpn_feature_map, backbone_feature_map, self.fuse_op) new_fpn_feature_map = fpn_block(fused_feature_map) fpn_feature_maps.append(new_fpn_feature_map) fpn_feature_maps.reverse() # bottom-up PAN pan_feature_maps = [fpn_feature_maps[0]] for idx, (downsample_conv, pan_block) in enumerate(zip(self.downsample_convs, self.pan_blocks)): top_pan_feature_map = pan_feature_maps[-1] fpn_feature_map = fpn_feature_maps[idx + 1] downsampled_feature_map = downsample_conv(top_pan_feature_map) fused_feature_map = fuse_feature_maps(downsampled_feature_map, fpn_feature_map, self.fuse_op) new_pan_feature_map = pan_block(fused_feature_map) pan_feature_maps.append(new_pan_feature_map) return Deimv2EncoderOutput(feature_maps=pan_feature_maps) 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) def weighting_function(max_num_bins: int, up: torch.Tensor, reg_scale: int) -> torch.Tensor: """ Generates the non-uniform Weighting Function W(n) for bounding box regression. Args: max_num_bins (int): Max number of the discrete bins. up (Tensor): Controls upper bounds of the sequence, where maximum offset is ±up * H / W. reg_scale (float): Controls the curvature of the Weighting Function. Larger values result in flatter weights near the central axis W(max_num_bins/2)=0 and steeper weights at both ends. Returns: Tensor: Sequence of Weighting Function. """ upper_bound1 = abs(up[0]) * abs(reg_scale) upper_bound2 = abs(up[0]) * abs(reg_scale) * 2 step = (upper_bound1 + 1) ** (2 / (max_num_bins - 2)) left_values = [-((step) ** i) + 1 for i in range(max_num_bins // 2 - 1, 0, -1)] right_values = [(step) ** i - 1 for i in range(1, max_num_bins // 2)] values = [-upper_bound2] + left_values + [torch.zeros_like(up[0][None])] + right_values + [upper_bound2] values = torch.cat(values, 0) return values def distance2bbox(points, distance: torch.Tensor, reg_scale: float) -> torch.Tensor: """ Decodes edge-distances into bounding box coordinates. Args: points (`torch.Tensor`): (batch_size, num_boxes, 4) or (num_boxes, 4) format, representing [x_center, y_center, width, height] distance (`torch.Tensor`): (batch_size, num_boxes, 4) or (num_boxes, 4), representing distances from the point to the left, top, right, and bottom boundaries. reg_scale (`float`): Controls the curvature of the Weighting Function. Returns: `torch.Tensor`: Bounding boxes in (batch_size, num_boxes, 4) or (num_boxes, 4) format, representing [x_center, y_center, width, height] """ reg_scale = abs(reg_scale) top_left_x = points[..., 0] - (0.5 * reg_scale + distance[..., 0]) * (points[..., 2] / reg_scale) top_left_y = points[..., 1] - (0.5 * reg_scale + distance[..., 1]) * (points[..., 3] / reg_scale) bottom_right_x = points[..., 0] + (0.5 * reg_scale + distance[..., 2]) * (points[..., 2] / reg_scale) bottom_right_y = points[..., 1] + (0.5 * reg_scale + distance[..., 3]) * (points[..., 3] / reg_scale) bboxes = torch.stack([top_left_x, top_left_y, bottom_right_x, bottom_right_y], -1) return corners_to_center_format(bboxes) class Deimv2Decoder(Deimv2PreTrainedModel): """ D-FINE Decoder implementing Fine-grained Distribution Refinement (FDR). This decoder refines object detection predictions through iterative updates across multiple layers, utilizing attention mechanisms, location quality estimators, and distribution refinement techniques to improve bounding box accuracy and robustness. """ _can_record_outputs = { "hidden_states": Deimv2DecoderLayer, "attentions": Deimv2SelfAttention, "cross_attentions": Deimv2MultiscaleDeformableAttention, } def __init__(self, config: Deimv2Config): super().__init__(config) self.eval_idx = config.eval_idx if config.eval_idx >= 0 else config.decoder_layers + config.eval_idx self.dropout = config.dropout self.layers = nn.ModuleList( [Deimv2DecoderLayer(config) for _ in range(config.decoder_layers)] + [Deimv2DecoderLayer(config) for _ in range(config.decoder_layers - self.eval_idx - 1)] ) self.query_pos_head = Deimv2MLP(4, config.d_model, config.d_model, 3, config.decoder_activation_function) # hack implementation for iterative bounding box refinement and two-stage Deformable DETR self.bbox_embed = None self.class_embed = None self.reg_scale = nn.Parameter(torch.tensor([config.reg_scale]), requires_grad=False) self.max_num_bins = config.max_num_bins self.d_model = config.d_model self.layer_scale = config.layer_scale self.pre_bbox_head = Deimv2MLP(config.hidden_size, config.hidden_size, 4, 3) self.integral = Deimv2Integral(config) self.num_head = config.decoder_attention_heads self.up = nn.Parameter(torch.tensor([config.up]), requires_grad=False) self.lqe_layers = nn.ModuleList([Deimv2LQE(config) for _ in range(config.decoder_layers)]) # Initialize weights and apply final processing self.post_init() @merge_with_config_defaults @capture_outputs def forward( self, encoder_hidden_states: torch.Tensor, reference_points: torch.Tensor, inputs_embeds: torch.Tensor, spatial_shapes, level_start_index=None, spatial_shapes_list=None, encoder_attention_mask=None, memory_mask=None, **kwargs: Unpack[TransformersKwargs], ) -> Deimv2DecoderOutput: r""" Args: inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`): The query embeddings that are passed into the decoder. encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, 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, 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**). reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)` is `as_two_stage` else `(batch_size, num_queries, 2)` or , *optional*): Reference point in range `[0, 1]`, top-left (0,0), bottom-right (1, 1), including padding area. spatial_shapes (`torch.FloatTensor` of shape `(num_feature_levels, 2)`): Spatial shapes of the feature maps. level_start_index (`torch.LongTensor` of shape `(num_feature_levels)`, *optional*): Indexes for the start of each feature level. In range `[0, sequence_length]`. """ if inputs_embeds is not None: hidden_states = inputs_embeds # decoder layers intermediate = () intermediate_reference_points = () intermediate_logits = () intermediate_predicted_corners = () initial_reference_points = () output_detach = pred_corners_undetach = 0 project = weighting_function(self.max_num_bins, self.up, self.reg_scale) ref_points_detach = F.sigmoid(reference_points) for i, decoder_layer in enumerate(self.layers): ref_points_input = ref_points_detach.unsqueeze(2) query_pos_embed = self.query_pos_head(ref_points_detach).clamp(min=-10, max=10) hidden_states = decoder_layer( hidden_states, position_embeddings=query_pos_embed, reference_points=ref_points_input, spatial_shapes=spatial_shapes, spatial_shapes_list=spatial_shapes_list, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, **kwargs, ) if i == 0: # Initial bounding box predictions with inverse sigmoid refinement new_reference_points = F.sigmoid( self.pre_bbox_head(hidden_states) + inverse_sigmoid(ref_points_detach) ) ref_points_initial = new_reference_points.detach() # Refine bounding box corners using FDR, integrating previous layer's corrections if self.bbox_embed is not None: pred_corners = self.bbox_embed[i](hidden_states + output_detach) + pred_corners_undetach inter_ref_bbox = distance2bbox( ref_points_initial, self.integral(pred_corners, project), self.reg_scale ) pred_corners_undetach = pred_corners ref_points_detach = inter_ref_bbox.detach() output_detach = hidden_states.detach() intermediate += (hidden_states,) if self.class_embed is not None and (self.training or i == self.eval_idx): scores = self.class_embed[i](hidden_states) # Add initial logits and reference points with pre-bbox head if i == 0: intermediate_logits += (scores,) intermediate_reference_points += (new_reference_points,) # Lqe does not affect the performance here. scores = self.lqe_layers[i](scores, pred_corners) intermediate_logits += (scores,) intermediate_reference_points += (inter_ref_bbox,) initial_reference_points += (ref_points_initial,) intermediate_predicted_corners += (pred_corners,) # Keep batch_size as first dimension intermediate = torch.stack(intermediate) if self.class_embed is not None and self.bbox_embed is not None: intermediate_logits = torch.stack(intermediate_logits, dim=1) intermediate_predicted_corners = torch.stack(intermediate_predicted_corners, dim=1) initial_reference_points = torch.stack(initial_reference_points, dim=1) intermediate_reference_points = torch.stack(intermediate_reference_points, dim=1) return Deimv2DecoderOutput( last_hidden_state=hidden_states, intermediate_hidden_states=intermediate, intermediate_logits=intermediate_logits, intermediate_reference_points=intermediate_reference_points, intermediate_predicted_corners=intermediate_predicted_corners, initial_reference_points=initial_reference_points, ) def get_contrastive_denoising_training_group( targets, num_classes, num_queries, class_embed, num_denoising_queries=100, label_noise_ratio=0.5, box_noise_scale=1.0, ): """ Creates a contrastive denoising training group using ground-truth samples. It adds noise to labels and boxes. Args: targets (`list[dict]`): The target objects, each containing 'class_labels' and 'boxes' for objects in an image. num_classes (`int`): Total number of classes in the dataset. num_queries (`int`): Number of query slots in the transformer. class_embed (`callable`): A function or a model layer to embed class labels. num_denoising_queries (`int`, *optional*, defaults to 100): Number of denoising queries. label_noise_ratio (`float`, *optional*, defaults to 0.5): Ratio of noise applied to labels. box_noise_scale (`float`, *optional*, defaults to 1.0): Scale of noise applied to bounding boxes. Returns: `tuple` comprising various elements: - **input_query_class** (`torch.FloatTensor`) -- Class queries with applied label noise. - **input_query_bbox** (`torch.FloatTensor`) -- Bounding box queries with applied box noise. - **attn_mask** (`torch.FloatTensor`) -- Attention mask for separating denoising and reconstruction queries. - **denoising_meta_values** (`dict`) -- Metadata including denoising positive indices, number of groups, and split sizes. """ if num_denoising_queries <= 0: return None, None, None, None num_ground_truths = [len(t["class_labels"]) for t in targets] device = targets[0]["class_labels"].device max_gt_num = max(num_ground_truths) if max_gt_num == 0: return None, None, None, None num_groups_denoising_queries = num_denoising_queries // max_gt_num num_groups_denoising_queries = 1 if num_groups_denoising_queries == 0 else num_groups_denoising_queries # pad gt to max_num of a batch batch_size = len(num_ground_truths) input_query_class = torch.full([batch_size, max_gt_num], num_classes, dtype=torch.int32, device=device) input_query_bbox = torch.zeros([batch_size, max_gt_num, 4], device=device) pad_gt_mask = torch.zeros([batch_size, max_gt_num], dtype=torch.bool, device=device) for i in range(batch_size): num_gt = num_ground_truths[i] if num_gt > 0: input_query_class[i, :num_gt] = targets[i]["class_labels"] input_query_bbox[i, :num_gt] = targets[i]["boxes"] pad_gt_mask[i, :num_gt] = 1 # each group has positive and negative queries. input_query_class = input_query_class.tile([1, 2 * num_groups_denoising_queries]) input_query_bbox = input_query_bbox.tile([1, 2 * num_groups_denoising_queries, 1]) pad_gt_mask = pad_gt_mask.tile([1, 2 * num_groups_denoising_queries]) # positive and negative mask negative_gt_mask = torch.zeros([batch_size, max_gt_num * 2, 1], device=device) negative_gt_mask[:, max_gt_num:] = 1 negative_gt_mask = negative_gt_mask.tile([1, num_groups_denoising_queries, 1]) positive_gt_mask = 1 - negative_gt_mask # contrastive denoising training positive index positive_gt_mask = positive_gt_mask.squeeze(-1) * pad_gt_mask denoise_positive_idx = torch.nonzero(positive_gt_mask)[:, 1] denoise_positive_idx = torch.split( denoise_positive_idx, [n * num_groups_denoising_queries for n in num_ground_truths] ) # total denoising queries num_denoising_queries = torch_int(max_gt_num * 2 * num_groups_denoising_queries) if label_noise_ratio > 0: mask = torch.rand_like(input_query_class, dtype=torch.float) < (label_noise_ratio * 0.5) # randomly put a new one here new_label = torch.randint_like(mask, 0, num_classes, dtype=input_query_class.dtype) input_query_class = torch.where(mask & pad_gt_mask, new_label, input_query_class) if box_noise_scale > 0: known_bbox = center_to_corners_format(input_query_bbox) diff = torch.tile(input_query_bbox[..., 2:] * 0.5, [1, 1, 2]) * box_noise_scale rand_sign = torch.randint_like(input_query_bbox, 0, 2) * 2.0 - 1.0 rand_part = torch.rand_like(input_query_bbox) rand_part = (rand_part + 1.0) * negative_gt_mask + rand_part * (1 - negative_gt_mask) rand_part *= rand_sign known_bbox += rand_part * diff known_bbox.clip_(min=0.0, max=1.0) input_query_bbox = corners_to_center_format(known_bbox) input_query_bbox = inverse_sigmoid(input_query_bbox) input_query_class = class_embed(input_query_class) target_size = num_denoising_queries + num_queries attn_mask = torch.full([target_size, target_size], 0, dtype=torch.float, device=device) # match query cannot see the reconstruction attn_mask[num_denoising_queries:, :num_denoising_queries] = -torch.inf # reconstructions cannot see each other for i in range(num_groups_denoising_queries): idx_block_start = max_gt_num * 2 * i idx_block_end = max_gt_num * 2 * (i + 1) attn_mask[idx_block_start:idx_block_end, :idx_block_start] = -torch.inf attn_mask[idx_block_start:idx_block_end, idx_block_end:num_denoising_queries] = -torch.inf denoising_meta_values = { "dn_positive_idx": denoise_positive_idx, "dn_num_group": num_groups_denoising_queries, "dn_num_split": [num_denoising_queries, num_queries], } return input_query_class, input_query_bbox, attn_mask, denoising_meta_values @auto_docstring( custom_intro=""" RT-DETR Model (consisting of a backbone and encoder-decoder) outputting raw hidden states without any head on top. """ ) class Deimv2Model(Deimv2PreTrainedModel): def __init__(self, config: Deimv2Config): super().__init__(config) is_dinov3 = getattr(config.backbone_config, "model_type", None) == "dinov3_vit" self.conv_encoder = Deimv2DINOv3ConvEncoder(config) if is_dinov3 else Deimv2ConvEncoder(config) self.encoder = ( Deimv2LiteEncoder(config) if config.encoder_type == "lite" else Deimv2HybridEncoder(config=config) ) if config.num_denoising > 0: self.denoising_class_embed = nn.Embedding( config.num_labels + 1, config.d_model, padding_idx=config.num_labels ) if config.learn_initial_query: self.weight_embedding = nn.Embedding(config.num_queries, config.d_model) self.enc_output = nn.Sequential( nn.Linear(config.d_model, config.d_model), nn.LayerNorm(config.d_model, eps=config.layer_norm_eps), ) self.enc_score_head = nn.Linear(config.d_model, config.num_labels) self.enc_bbox_head = Deimv2MLP(config.d_model, config.d_model, 4, 3) if config.anchor_image_size: self.anchors, self.valid_mask = self.generate_anchors(dtype=self.dtype) num_backbone_outs = len(config.decoder_in_channels) decoder_input_proj = [] in_channels = config.decoder_in_channels[-1] for _ in range(num_backbone_outs): decoder_input_proj.append( nn.Identity() if config.hidden_size == config.decoder_in_channels[-1] else Deimv2ConvNormLayer(config, in_channels, config.d_model, 1, 1) ) for _ in range(config.num_feature_levels - num_backbone_outs): decoder_input_proj.append( nn.Identity() if config.hidden_size == config.decoder_in_channels[-1] else Deimv2ConvNormLayer(config, in_channels, config.d_model, 3, 2) ) self.decoder_input_proj = nn.ModuleList(decoder_input_proj) self.decoder = Deimv2Decoder(config) self.post_init() def freeze_backbone(self): for param in self.backbone.parameters(): param.requires_grad_(False) def unfreeze_backbone(self): for param in self.backbone.parameters(): param.requires_grad_(True) @compile_compatible_method_lru_cache(maxsize=32) def generate_anchors(self, spatial_shapes=None, grid_size=0.05, device="cpu", dtype=torch.float32): if spatial_shapes is None: spatial_shapes = [ [int(self.config.anchor_image_size[0] / s), int(self.config.anchor_image_size[1] / s)] for s in self.config.feat_strides ] anchors = [] for level, (height, width) in enumerate(spatial_shapes): grid_y, grid_x = torch.meshgrid( torch.arange(end=height, device=device).to(dtype), torch.arange(end=width, device=device).to(dtype), indexing="ij", ) grid_xy = torch.stack([grid_x, grid_y], -1) grid_xy = grid_xy.unsqueeze(0) + 0.5 grid_xy[..., 0] /= width grid_xy[..., 1] /= height wh = torch.ones_like(grid_xy) * grid_size * (2.0**level) anchors.append(torch.concat([grid_xy, wh], -1).reshape(-1, height * width, 4)) # define the valid range for anchor coordinates eps = 1e-2 anchors = torch.concat(anchors, 1) valid_mask = ((anchors > eps) * (anchors < 1 - eps)).all(-1, keepdim=True) anchors = torch.log(anchors / (1 - anchors)) anchors = torch.where(valid_mask, anchors, torch.tensor(torch.finfo(dtype).max, dtype=dtype, device=device)) return anchors, valid_mask @auto_docstring @can_return_tuple def forward( self, pixel_values: torch.FloatTensor, pixel_mask: torch.LongTensor | None = None, encoder_outputs: torch.FloatTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: list[dict] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor] | Deimv2ModelOutput: r""" 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. 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, Deimv2Model >>> 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("PekingU/Deimv2_r50vd") >>> model = Deimv2Model.from_pretrained("PekingU/Deimv2_r50vd") >>> inputs = image_processor(images=image, return_tensors="pt") >>> outputs = model(**inputs) >>> last_hidden_states = outputs.last_hidden_state >>> list(last_hidden_states.shape) [1, 300, 256] ```""" # Overrides DFineModel.forward: DEIMv2 uses a unified conv_encoder (backbone + projection) instead of # D-FINE's separate backbone + encoder_input_proj, and returns feature_maps instead of last_hidden_state. if pixel_values is None and inputs_embeds is None: raise ValueError("You have to specify either pixel_values or inputs_embeds") # extract multi-scale features via conv_encoder (backbone + projection in one step) if inputs_embeds is None: 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) # TODO: pass pixel_mask to backbone once DINOv3 supports it proj_feats = self.conv_encoder(pixel_values) else: batch_size = inputs_embeds.shape[0] device = inputs_embeds.device proj_feats = inputs_embeds encoder_outputs = self.encoder( proj_feats, **kwargs, ) # Equivalent to def _get_encoder_input # https://github.com/lyuwenyu/RT-DETR/blob/94f5e16708329d2f2716426868ec89aa774af016/rtdetr_pytorch/src/zoo/rtdetr/rtdetr_decoder.py#L412 sources = [] for level, source in enumerate(encoder_outputs.feature_maps): sources.append(self.decoder_input_proj[level](source)) # Lowest resolution feature maps are obtained via 3x3 stride 2 convolutions on the final stage if self.config.num_feature_levels > len(sources): sources.append(self.decoder_input_proj[len(sources)](encoder_outputs.feature_maps[-1])) for i in range(len(sources), self.config.num_feature_levels): sources.append(self.decoder_input_proj[i](encoder_outputs.feature_maps[-1])) # Prepare encoder inputs (by flattening) source_flatten = [] spatial_shapes_list = [] spatial_shapes = torch.empty((len(sources), 2), device=device, dtype=torch.long) for level, source in enumerate(sources): height, width = source.shape[-2:] spatial_shapes[level, 0] = height spatial_shapes[level, 1] = width spatial_shapes_list.append((height, width)) source = source.flatten(2).transpose(1, 2) source_flatten.append(source) source_flatten = torch.cat(source_flatten, 1) level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1])) # prepare denoising training if self.training and self.config.num_denoising > 0 and labels is not None: ( denoising_class, denoising_bbox_unact, attention_mask, denoising_meta_values, ) = get_contrastive_denoising_training_group( targets=labels, num_classes=self.config.num_labels, num_queries=self.config.num_queries, class_embed=self.denoising_class_embed, num_denoising_queries=self.config.num_denoising, label_noise_ratio=self.config.label_noise_ratio, box_noise_scale=self.config.box_noise_scale, ) else: denoising_class, denoising_bbox_unact, attention_mask, denoising_meta_values = None, None, None, None batch_size = len(source_flatten) device = source_flatten.device dtype = source_flatten.dtype # prepare input for decoder if self.training or self.config.anchor_image_size is None: # Pass spatial_shapes as tuple to make it hashable and make sure # lru_cache is working for generate_anchors() spatial_shapes_tuple = tuple(spatial_shapes_list) anchors, valid_mask = self.generate_anchors(spatial_shapes_tuple, device=device, dtype=dtype) else: anchors, valid_mask = self.anchors, self.valid_mask anchors, valid_mask = anchors.to(device, dtype), valid_mask.to(device, dtype) # use the valid_mask to selectively retain values in the feature map where the mask is True memory = valid_mask.to(source_flatten.dtype) * source_flatten output_memory = self.enc_output(memory) enc_outputs_class = self.enc_score_head(output_memory) enc_outputs_coord_logits = self.enc_bbox_head(output_memory) + anchors _, topk_ind = torch.topk(enc_outputs_class.max(-1).values, self.config.num_queries, dim=1) reference_points_unact = enc_outputs_coord_logits.gather( dim=1, index=topk_ind.unsqueeze(-1).repeat(1, 1, enc_outputs_coord_logits.shape[-1]) ) enc_topk_bboxes = F.sigmoid(reference_points_unact) if denoising_bbox_unact is not None: reference_points_unact = torch.concat([denoising_bbox_unact, reference_points_unact], 1) enc_topk_logits = enc_outputs_class.gather( dim=1, index=topk_ind.unsqueeze(-1).repeat(1, 1, enc_outputs_class.shape[-1]) ) # extract region features if self.config.learn_initial_query: target = self.weight_embedding.tile([batch_size, 1, 1]) else: target = output_memory.gather(dim=1, index=topk_ind.unsqueeze(-1).repeat(1, 1, output_memory.shape[-1])) target = target.detach() if denoising_class is not None: target = torch.concat([denoising_class, target], 1) init_reference_points = reference_points_unact.detach() # decoder decoder_outputs = self.decoder( inputs_embeds=target, encoder_hidden_states=source_flatten, encoder_attention_mask=attention_mask, reference_points=init_reference_points, spatial_shapes=spatial_shapes, spatial_shapes_list=spatial_shapes_list, level_start_index=level_start_index, **kwargs, ) return Deimv2ModelOutput( last_hidden_state=decoder_outputs.last_hidden_state, intermediate_hidden_states=decoder_outputs.intermediate_hidden_states, intermediate_logits=decoder_outputs.intermediate_logits, intermediate_reference_points=decoder_outputs.intermediate_reference_points, intermediate_predicted_corners=decoder_outputs.intermediate_predicted_corners, initial_reference_points=decoder_outputs.initial_reference_points, decoder_hidden_states=decoder_outputs.hidden_states, decoder_attentions=decoder_outputs.attentions, cross_attentions=decoder_outputs.cross_attentions, encoder_last_hidden_state=encoder_outputs.feature_maps, encoder_hidden_states=encoder_outputs.hidden_states, encoder_attentions=encoder_outputs.attentions, init_reference_points=init_reference_points, enc_topk_logits=enc_topk_logits, enc_topk_bboxes=enc_topk_bboxes, enc_outputs_class=enc_outputs_class, enc_outputs_coord_logits=enc_outputs_coord_logits, denoising_meta_values=denoising_meta_values, ) @dataclass @auto_docstring( custom_intro=""" Output type of [`Deimv2ForObjectDetection`]. """ ) class Deimv2ObjectDetectionOutput(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 [`~Deimv2ImageProcessor.post_process_object_detection`] to retrieve the unnormalized (absolute) 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, num_queries, 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 `(batch_size, config.decoder_layers, num_queries, hidden_size)`): Stacked intermediate hidden states (output of each layer of the decoder). intermediate_logits (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, config.num_labels)`): Stacked intermediate logits (logits of each layer of the decoder). intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): Stacked intermediate reference points (reference points of each layer of the decoder). intermediate_predicted_corners (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): Stacked intermediate predicted corners (predicted corners of each layer of the decoder). initial_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`): Stacked initial reference points (initial reference points of each layer of the decoder). init_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`): Initial reference points sent through the Transformer decoder. enc_topk_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): Logits of predicted bounding boxes coordinates in the encoder. enc_topk_bboxes (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): Logits of predicted bounding boxes coordinates in the encoder. enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are picked as region proposals in the first stage. Output of bounding box binary classification (i.e. foreground and background). enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`): Logits of predicted bounding boxes coordinates in the first stage. denoising_meta_values (`dict`): Extra dictionary for the denoising related values """ 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 intermediate_hidden_states: torch.FloatTensor | None = None intermediate_logits: torch.FloatTensor | None = None intermediate_reference_points: torch.FloatTensor | None = None intermediate_predicted_corners: torch.FloatTensor | None = None initial_reference_points: 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 init_reference_points: tuple[torch.FloatTensor] | None = None enc_topk_logits: torch.FloatTensor | None = None enc_topk_bboxes: torch.FloatTensor | None = None enc_outputs_class: torch.FloatTensor | None = None enc_outputs_coord_logits: torch.FloatTensor | None = None denoising_meta_values: dict | None = None @auto_docstring( custom_intro=""" RT-DETR Model (consisting of a backbone and encoder-decoder) outputting bounding boxes and logits to be further decoded into scores and classes. """ ) class Deimv2ForObjectDetection(Deimv2PreTrainedModel): _tied_weights_keys = { r"bbox_embed.(?![0])\d+": r"bbox_embed.0", r"class_embed.(?![0])\d+": r"^class_embed.0", "class_embed": "model.decoder.class_embed", "bbox_embed": "model.decoder.bbox_embed", } def __init__(self, config: Deimv2Config): super().__init__(config) self.eval_idx = config.eval_idx if config.eval_idx >= 0 else config.decoder_layers + config.eval_idx self.model = Deimv2Model(config) scaled_dim = round(config.layer_scale * config.hidden_size) num_pred = config.decoder_layers self.class_embed = nn.ModuleList([nn.Linear(config.d_model, config.num_labels) for _ in range(num_pred)]) if config.share_bbox_head: shared_bbox = Deimv2MLP(config.hidden_size, config.hidden_size, 4 * (config.max_num_bins + 1), 3) self.bbox_embed = nn.ModuleList([shared_bbox] * num_pred) else: self.bbox_embed = nn.ModuleList( [ Deimv2MLP(config.hidden_size, config.hidden_size, 4 * (config.max_num_bins + 1), 3) for _ in range(self.eval_idx + 1) ] + [ Deimv2MLP(scaled_dim, scaled_dim, 4 * (config.max_num_bins + 1), 3) for _ in range(config.decoder_layers - self.eval_idx - 1) ] ) self.model.decoder.class_embed = self.class_embed self.model.decoder.bbox_embed = self.bbox_embed self.post_init() def _set_aux_loss(self, outputs_class, outputs_coord): return [{"logits": a, "pred_boxes": b} for a, b in zip(outputs_class, outputs_coord)] @auto_docstring @can_return_tuple def forward( self, pixel_values: torch.FloatTensor, pixel_mask: torch.LongTensor | None = None, encoder_outputs: torch.FloatTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: list[dict] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor] | Deimv2ObjectDetectionOutput: r""" Example: ```python >>> import torch >>> from transformers.image_utils import load_image >>> from transformers import AutoImageProcessor, Deimv2ForObjectDetection >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = load_image(url) >>> image_processor = AutoImageProcessor.from_pretrained("harshaljanjani/DEIMv2_HGNetv2_N_COCO_Transformers") >>> model = Deimv2ForObjectDetection.from_pretrained("harshaljanjani/DEIMv2_HGNetv2_N_COCO_Transformers") >>> # prepare image for the model >>> inputs = image_processor(images=image, return_tensors="pt") >>> # forward pass >>> outputs = model(**inputs) >>> logits = outputs.logits >>> list(logits.shape) [1, 300, 80] >>> boxes = outputs.pred_boxes >>> list(boxes.shape) [1, 300, 4] >>> # 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.9, target_sizes=target_sizes) >>> result = results[0] # first image in batch >>> for score, label, box in zip(result["scores"], result["labels"], result["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}" ... ) ``` """ outputs = self.model( pixel_values, pixel_mask=pixel_mask, encoder_outputs=encoder_outputs, inputs_embeds=inputs_embeds, labels=labels, **kwargs, ) denoising_meta_values = outputs.denoising_meta_values if self.training else None outputs_class = outputs.intermediate_logits outputs_coord = outputs.intermediate_reference_points predicted_corners = outputs.intermediate_predicted_corners initial_reference_points = outputs.initial_reference_points logits = outputs_class[:, -1] pred_boxes = outputs_coord[:, -1] loss, loss_dict, auxiliary_outputs, enc_topk_logits, enc_topk_bboxes = None, None, None, None, None if labels is not None: enc_topk_logits = outputs.enc_topk_logits enc_topk_bboxes = outputs.enc_topk_bboxes loss, loss_dict, auxiliary_outputs = self.loss_function( logits, labels, self.device, pred_boxes, self.config, outputs_class, outputs_coord, enc_topk_logits=enc_topk_logits, enc_topk_bboxes=enc_topk_bboxes, denoising_meta_values=denoising_meta_values, predicted_corners=predicted_corners, initial_reference_points=initial_reference_points, **kwargs, ) return Deimv2ObjectDetectionOutput( loss=loss, loss_dict=loss_dict, logits=logits, pred_boxes=pred_boxes, auxiliary_outputs=auxiliary_outputs, last_hidden_state=outputs.last_hidden_state, intermediate_hidden_states=outputs.intermediate_hidden_states, intermediate_logits=outputs.intermediate_logits, intermediate_reference_points=outputs.intermediate_reference_points, intermediate_predicted_corners=outputs.intermediate_predicted_corners, initial_reference_points=outputs.initial_reference_points, 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, init_reference_points=outputs.init_reference_points, enc_topk_logits=outputs.enc_topk_logits, enc_topk_bboxes=outputs.enc_topk_bboxes, enc_outputs_class=outputs.enc_outputs_class, enc_outputs_coord_logits=outputs.enc_outputs_coord_logits, denoising_meta_values=outputs.denoising_meta_values, ) @property def _tied_weights_keys(self): keys = { r"class_embed.(?![0])\d+": r"^class_embed.0", "class_embed": "model.decoder.class_embed", "bbox_embed": "model.decoder.bbox_embed", } if self.config.share_bbox_head: keys[r"model\.decoder\.bbox_embed\.(?![0])\d+"] = r"model.decoder.bbox_embed.0" keys[r"bbox_embed.(?![0])\d+"] = r"bbox_embed.0" return keys __all__ = ["Deimv2Model", "Deimv2PreTrainedModel", "Deimv2ForObjectDetection"]