# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/lw_detr/modular_lw_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_lw_detr.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 collections.abc import math import warnings from collections.abc import Callable from dataclasses import dataclass from typing import Any import torch import torch.nn.functional as F from torch import Tensor, nn from ... import initialization as init from ...activations import ACT2CLS, ACT2FN from ...backbone_utils import BackboneMixin from ...integrations import use_kernel_forward_from_hub from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BackboneOutput, BaseModelOutputWithCrossAttentions from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import ModelOutput, TransformersKwargs, auto_docstring, torch_compilable_check from ...utils.generic import can_return_tuple, merge_with_config_defaults from ...utils.output_capturing import capture_outputs from .configuration_lw_detr import LwDetrConfig, LwDetrViTConfig def eager_attention_forward( module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: torch.Tensor | None, scaling: float, dropout: float = 0.0, **kwargs: Unpack[TransformersKwargs], ): key_states = repeat_kv(key, module.num_key_value_groups) value_states = repeat_kv(value, module.num_key_value_groups) attn_weights = torch.matmul(query, key_states.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, dtype=torch.float32).to(query.dtype) attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() return attn_output, attn_weights def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) """ batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) class LwDetrViTSelfAttention(nn.Module): def __init__(self, config: LwDetrViTConfig): super().__init__() if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"): raise ValueError( f"The hidden size {config.hidden_size} is not a multiple of the number of attention " f"heads {config.num_attention_heads}." ) self.config = config self.num_attention_heads = config.num_attention_heads self.attention_head_size = int(config.hidden_size / config.num_attention_heads) self.all_head_size = self.num_attention_heads * self.attention_head_size self.dropout_prob = config.dropout_prob self.scaling = self.attention_head_size**-0.5 self.is_causal = False self.query = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias) self.value = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias) self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=False) self.num_key_value_groups = 1 def forward( self, hidden_states: torch.Tensor, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor, torch.Tensor]: batch_size = hidden_states.shape[0] new_shape = batch_size, -1, self.num_attention_heads, self.attention_head_size key_layer = self.key(hidden_states).view(*new_shape).transpose(1, 2) value_layer = self.value(hidden_states).view(*new_shape).transpose(1, 2) query_layer = self.query(hidden_states).view(*new_shape).transpose(1, 2) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward ) context_layer, attention_probs = attention_interface( self, query_layer, key_layer, value_layer, None, is_causal=self.is_causal, scaling=self.scaling, dropout=0.0 if not self.training else self.dropout_prob, **kwargs, ) new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) context_layer = context_layer.reshape(new_context_layer_shape) return context_layer, attention_probs class LwDetrViTAttention(nn.Module): def __init__(self, config: LwDetrViTConfig): """ Args: config (`LwDetrViTConfig`): Model configuration. """ super().__init__() self.attention = LwDetrViTSelfAttention(config) self.output = nn.Linear(config.hidden_size, config.hidden_size) def forward( self, hidden_states: torch.Tensor, **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: self_attn_output, _ = self.attention(hidden_states, **kwargs) output = self.output(self_attn_output) return output class LwDetrViTMlp(nn.Module): def __init__(self, config, in_features: int, hidden_features: int) -> None: super().__init__() self.fc1 = nn.Linear(in_features, hidden_features) self.act = ACT2FN[config.hidden_act] self.fc2 = nn.Linear(hidden_features, in_features) self.drop = nn.Dropout(config.dropout_prob) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.fc1(x) x = self.act(x) x = self.drop(x) x = self.fc2(x) x = self.drop(x) return x class LwDetrViTLayer(GradientCheckpointingLayer): def __init__( self, config: LwDetrViTConfig, layer_idx, ) -> None: super().__init__() dim = config.hidden_size self.attention = LwDetrViTAttention(config) self.intermediate = LwDetrViTMlp(config=config, in_features=dim, hidden_features=int(dim * config.mlp_ratio)) self.layernorm_before = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.layernorm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.gamma_1 = nn.Parameter(torch.Tensor(dim), requires_grad=True) self.gamma_2 = nn.Parameter(torch.Tensor(dim), requires_grad=True) self.window = layer_idx in config.window_block_indices self.num_windows = config.num_windows def forward( self, hidden_states: torch.Tensor, **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: batch_size, seq_len, channels = hidden_states.shape hidden_states_norm = self.layernorm_before(hidden_states) if not self.window: hidden_states_norm = hidden_states_norm.reshape( batch_size // self.num_windows, self.num_windows * seq_len, channels ) attention_output = self.attention(hidden_states_norm, **kwargs) attention_output = attention_output * self.gamma_1 if not self.window: attention_output = attention_output.reshape(batch_size, seq_len, channels) hidden_states = hidden_states + attention_output layer_output = self.layernorm_after(hidden_states) layer_output = self.intermediate(layer_output) layer_output = layer_output * self.gamma_2 hidden_states = hidden_states + layer_output return hidden_states class LwDetrViTEncoder(nn.Module): def __init__(self, config: LwDetrViTConfig) -> None: super().__init__() self.config = config self.layer = nn.ModuleList([LwDetrViTLayer(config, i) for i in range(config.num_hidden_layers)]) self.gradient_checkpointing = False def forward( self, hidden_states: torch.Tensor, **kwargs: Unpack[TransformersKwargs], ) -> list[torch.Tensor]: list_hidden_states = [hidden_states] for i, layer_module in enumerate(self.layer): hidden_states = layer_module(hidden_states, **kwargs) list_hidden_states.append(hidden_states) return list_hidden_states class LwDetrViTEmbeddings(nn.Module): """ This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial `hidden_states` (patch embeddings) to be consumed by a Transformer. """ def __init__(self, config): super().__init__() image_size, patch_size = config.pretrain_image_size, config.patch_size num_channels, hidden_size = config.num_channels, config.hidden_size image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size) patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size) num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0]) self.image_size = image_size self.patch_size = patch_size self.num_channels = num_channels self.num_patches = num_patches if config.use_absolute_position_embeddings: # Initialize absolute positional embedding with pretrain image size. num_positions = num_patches + 1 self.position_embeddings = nn.Parameter(torch.zeros(1, num_positions, config.hidden_size)) else: self.position_embeddings = None self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size) def get_absolute_positions(self, abs_pos_embeddings, has_cls_token, height, width): """ Calculate absolute positional embeddings. If needed, resize embeddings and remove cls_token dimension for the original embeddings. Args: abs_pos_embeddings (`torch.Tensor`): Absolute positional embeddings with (1, num_position, num_channels). has_cls_token (`bool`): If true, has 1 embedding in abs_pos_embeddings for cls token. height (`int`): Height of input image tokens. width (`int`): Width of input image tokens. Returns: Absolute positional embeddings after processing with shape (1, height, width, num_channels) """ if has_cls_token: abs_pos_embeddings = abs_pos_embeddings[:, 1:] num_position = abs_pos_embeddings.shape[1] size = int(math.sqrt(num_position)) # This is a constant and can be recorded as such in the ONNX export. if size * size != num_position: raise ValueError("Absolute position embeddings must be a square number.") if torch.jit.is_tracing() or (size != height or size != width): # nn.functional.interpolate is a noop in case size == height and size == width - we need to always capture this path with jit.trace. new_abs_pos_embeddings = nn.functional.interpolate( abs_pos_embeddings.reshape(1, size, size, -1).permute(0, 3, 1, 2), size=(height, width), mode="bicubic", align_corners=False, ) return new_abs_pos_embeddings.permute(0, 2, 3, 1) else: return abs_pos_embeddings.reshape(1, height, width, -1) def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: num_channels = pixel_values.shape[1] if num_channels != self.num_channels: raise ValueError( "Make sure that the channel dimension of the pixel values match with the one set in the configuration." f" Expected {self.num_channels} but got {num_channels}." ) embeddings = self.projection(pixel_values) if self.position_embeddings is not None: # (batch_size, num_channels, height, width) -> (batch_size, height, width, num_channels) embeddings = embeddings.permute(0, 2, 3, 1) # add position embeddings embeddings = embeddings + self.get_absolute_positions( self.position_embeddings, True, embeddings.shape[1], embeddings.shape[2] ) # (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width) embeddings = embeddings.permute(0, 3, 1, 2) return embeddings @auto_docstring class LwDetrViTPreTrainedModel(PreTrainedModel): config: LwDetrViTConfig base_model_prefix = "lw_detr_vit" main_input_name = "pixel_values" input_modalities = ("image",) supports_gradient_checkpointing = True _no_split_modules = ["LwDetrViTEmbeddings", "LwDetrViTLayer"] _supports_sdpa = True _supports_flash_attn = True _supports_flex_attn = True _supports_attention_backend = True _can_record_outputs = { "hidden_states": LwDetrViTLayer, "attentions": LwDetrViTSelfAttention, } @torch.no_grad() def _init_weights(self, module) -> None: """Initialize the weights""" if isinstance(module, (nn.Linear, nn.Conv2d)): init.trunc_normal_(module.weight, mean=0.0, std=self.config.initializer_range) if module.bias is not None: init.zeros_(module.bias) elif isinstance(module, nn.LayerNorm): init.zeros_(module.bias) init.ones_(module.weight) elif isinstance(module, LwDetrViTEmbeddings): init.trunc_normal_(module.position_embeddings, mean=0.0, std=self.config.initializer_range) if isinstance(module, LwDetrViTLayer): nn.init.constant_(module.gamma_1, self.config.cae_init_values) nn.init.constant_(module.gamma_2, self.config.cae_init_values) @auto_docstring() class LwDetrViTBackbone(BackboneMixin, LwDetrViTPreTrainedModel): def __init__(self, config): super().__init__(config) self.embeddings = LwDetrViTEmbeddings(config) self.encoder = LwDetrViTEncoder(config) self.num_features = [config.hidden_size for _ in range(config.num_hidden_layers + 1)] # initialize weights and apply final processing self.post_init() def get_input_embeddings(self) -> LwDetrViTEmbeddings: return self.embeddings.projection @merge_with_config_defaults @capture_outputs @auto_docstring def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BackboneOutput: r""" Examples: ```python >>> from transformers import LwDetrViTConfig, LwDetrViTBackbone >>> import torch >>> config = LwDetrViTConfig() >>> model = LwDetrViTBackbone(config) >>> pixel_values = torch.randn(1, 3, 224, 224) >>> with torch.no_grad(): ... outputs = model(pixel_values) >>> feature_maps = outputs.feature_maps >>> list(feature_maps[-1].shape) [1, 768, 14, 14] ```""" embedding_output = self.embeddings(pixel_values) batch_size, channels, height, width = embedding_output.shape # (batch_size, channels, height, width) -> (batch_size, height, width, channels) hidden_states = embedding_output.permute(0, 2, 3, 1) window_height = height // self.config.num_windows_side window_width = width // self.config.num_windows_side # (batch_size, height, width, channels) -> (batch_size*num_windows_side**2, window_height*window_width, channels) hidden_states = ( hidden_states.reshape( batch_size, self.config.num_windows_side, window_height, self.config.num_windows_side, window_width, channels, ) .permute(0, 1, 3, 2, 4, 5) .reshape(batch_size * self.config.num_windows_side**2, window_height * window_width, channels) ) hidden_states = self.encoder(hidden_states, **kwargs) feature_maps = () for stage, hidden_state in zip(self.stage_names, hidden_states): if stage in self.out_features: hidden_state = ( hidden_state.reshape( batch_size, self.config.num_windows_side, self.config.num_windows_side, window_height, window_width, channels, ) .permute(0, 5, 1, 3, 2, 4) .reshape(batch_size, channels, height, width) ) feature_maps += (hidden_state,) return BackboneOutput(feature_maps=feature_maps) class LwDetrConvNormLayer(nn.Module): def __init__( self, config: LwDetrConfig, in_channels: int, out_channels: int, kernel_size: int, stride: int, activation: str | None = None, ): super().__init__() self.conv = nn.Conv2d( in_channels, out_channels, kernel_size, stride, padding=kernel_size // 2, 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 LwDetrRepVggBlock(nn.Module): def __init__(self, config: LwDetrConfig): super().__init__() hidden_channels = int(config.d_model * config.hidden_expansion) self.conv1 = LwDetrConvNormLayer( config, hidden_channels, hidden_channels, 3, 1, activation=config.activation_function ) self.conv2 = LwDetrConvNormLayer( config, hidden_channels, hidden_channels, 3, 1, activation=config.activation_function ) def forward(self, x: torch.Tensor) -> torch.Tensor: y = self.conv1(x) y = self.conv2(y) return y class LwDetrC2FLayer(nn.Module): # Inspired by RTDetrCSPRepLayer def __init__(self, config: LwDetrConfig, in_channels: int): super().__init__() num_blocks = config.c2f_num_blocks activation = config.activation_function out_channels = config.d_model self.hidden_channels = int(out_channels * config.hidden_expansion) conv1_out_channels = 2 * self.hidden_channels self.conv1 = LwDetrConvNormLayer(config, in_channels, conv1_out_channels, 1, 1, activation=activation) conv2_in_channels = (2 + num_blocks) * self.hidden_channels self.conv2 = LwDetrConvNormLayer(config, conv2_in_channels, out_channels, 1, 1, activation=activation) self.bottlenecks = nn.ModuleList(LwDetrRepVggBlock(config) for _ in range(num_blocks)) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.conv1(hidden_states) all_hidden_states = list(hidden_states.split(self.hidden_channels, 1)) hidden_states = all_hidden_states[-1] for bottleneck in self.bottlenecks: hidden_states = bottleneck(hidden_states) all_hidden_states.append(hidden_states) hidden_states = torch.cat(all_hidden_states, 1) hidden_states = self.conv2(hidden_states) return hidden_states class LwDetrLayerNorm(nn.LayerNorm): r"""LayerNorm that supports two data formats: channels_last (default) or channels_first. The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height, width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width). """ def __init__(self, normalized_shape, *, eps=1e-6, data_format="channels_last", **kwargs): super().__init__(normalized_shape, eps=eps, **kwargs) if data_format not in ["channels_last", "channels_first"]: raise NotImplementedError(f"Unsupported data format: {data_format}") self.data_format = data_format def forward(self, features: torch.Tensor) -> torch.Tensor: """ Args: features: Tensor of shape (batch_size, channels, height, width) OR (batch_size, height, width, channels) """ if self.data_format == "channels_first": features = features.permute(0, 2, 3, 1) features = super().forward(features) features = features.permute(0, 3, 1, 2) else: features = super().forward(features) return features class LwDetrSamplingLayer(nn.Module): def __init__(self, config: LwDetrConfig, channel_size: int, scale: float): super().__init__() self.scale = scale self.channel_size = channel_size layers = [] if scale == 2.0: if channel_size > 512: layers.append(LwDetrConvNormLayer(config, channel_size, channel_size // 2, 1, 1, activation="relu")) layers.append(nn.ConvTranspose2d(channel_size // 2, channel_size // 4, kernel_size=2, stride=2)) else: layers.append(nn.ConvTranspose2d(channel_size, channel_size // 2, 2, 2)) elif scale == 0.5: layers.append(LwDetrConvNormLayer(config, channel_size, channel_size, 3, 2, activation="relu")) self.layers = nn.ModuleList(layers) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: for layer in self.layers: hidden_states = layer(hidden_states) return hidden_states class LwDetrScaleProjector(nn.Module): def __init__(self, config: LwDetrConfig, scale: float): super().__init__() intermediate_dims = [config.backbone_config.hidden_size] * len(config.backbone_config.out_indices) sampling_layers = [] for channel_size in intermediate_dims: sampling_layers.append(LwDetrSamplingLayer(config, channel_size, scale)) self.sampling_layers = nn.ModuleList(sampling_layers) intermediate_dim = intermediate_dims[-1] if scale == 2.0: if intermediate_dim > 512: intermediate_dim = intermediate_dim // 4 else: intermediate_dim = intermediate_dim // 2 projector_input_dim = intermediate_dim * len(intermediate_dims) self.projector_layer = LwDetrC2FLayer(config, projector_input_dim) self.layer_norm = LwDetrLayerNorm(config.d_model, data_format="channels_first") def forward(self, hidden_states_tuple: tuple[torch.Tensor]) -> torch.Tensor: sampled_hidden_states = [] for sampling_layer, hidden_states in zip(self.sampling_layers, hidden_states_tuple): hidden_states = sampling_layer(hidden_states) sampled_hidden_states.append(hidden_states) hidden_states = torch.cat(sampled_hidden_states, dim=1) hidden_states = self.projector_layer(hidden_states) hidden_states = self.layer_norm(hidden_states) return hidden_states class LwDetrMultiScaleProjector(nn.Module): def __init__(self, config: LwDetrConfig): super().__init__() self.config = config scale_factors = config.projector_scale_factors self.scale_layers = nn.ModuleList([LwDetrScaleProjector(config, scale) for scale in scale_factors]) def forward(self, hidden_states: tuple[torch.Tensor]) -> list[torch.Tensor]: output_hidden_states = [] for scale_layer in self.scale_layers: output_hidden_states.append(scale_layer(hidden_states)) return output_hidden_states class LwDetrConvEncoder(nn.Module): def __init__(self, config: LwDetrConfig): super().__init__() self.backbone = LwDetrViTBackbone(config.backbone_config) self.projector = LwDetrMultiScaleProjector(config) 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.backbone(pixel_values).feature_maps features = self.projector(features) 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 LwDetrAttention(nn.Module): def __init__(self, config: LwDetrConfig, layer_idx: int): super().__init__() self.config = config self.layer_idx = layer_idx self.head_dim = getattr(config, "head_dim", config.d_model // config.decoder_self_attention_heads) self.scaling = self.head_dim**-0.5 self.attention_dropout = config.attention_dropout self.is_causal = False self.num_key_value_groups = 1 self.q_proj = nn.Linear( config.d_model, config.decoder_self_attention_heads * self.head_dim, bias=config.attention_bias ) self.k_proj = nn.Linear( config.d_model, config.decoder_self_attention_heads * self.head_dim, bias=config.attention_bias ) self.v_proj = nn.Linear( config.d_model, config.decoder_self_attention_heads * self.head_dim, bias=config.attention_bias ) self.o_proj = nn.Linear( config.decoder_self_attention_heads * self.head_dim, config.d_model, bias=config.attention_bias ) def forward( self, hidden_states: torch.Tensor, position_embeddings: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor, torch.Tensor]: batch_size, seq_len, _ = hidden_states.shape input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) hidden_states_original = hidden_states if position_embeddings is not None: hidden_states = hidden_states if position_embeddings is None else hidden_states + position_embeddings if self.training: # at training, we use group detr technique to add more supervision by using multiple weight-sharing decoders at once for faster convergence # at inference, we only use one decoder hidden_states_original = torch.cat( hidden_states_original.split(seq_len // self.config.group_detr, dim=1), dim=0 ) hidden_states = torch.cat(hidden_states.split(seq_len // self.config.group_detr, dim=1), dim=0) query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) value_states = self.v_proj(hidden_states_original).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=None, 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) if self.training: attn_output = torch.cat(torch.split(attn_output, batch_size, dim=0), dim=1) return attn_output, attn_weights @use_kernel_forward_from_hub("MultiScaleDeformableAttention") class MultiScaleDeformableAttention(nn.Module): def forward( self, value: Tensor, value_spatial_shapes: Tensor, value_spatial_shapes_list: list[tuple], level_start_index: Tensor, sampling_locations: Tensor, attention_weights: Tensor, im2col_step: int, ): batch_size, _, num_heads, hidden_dim = value.shape _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape value_list = value.split([height * width for height, width in value_spatial_shapes_list], dim=1) sampling_grids = 2 * sampling_locations - 1 sampling_value_list = [] for level_id, (height, width) in enumerate(value_spatial_shapes_list): # 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] .flatten(2) .transpose(1, 2) .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].transpose(1, 2).flatten(0, 1) # batch_size*num_heads, hidden_dim, num_queries, num_points sampling_value_l_ = nn.functional.grid_sample( value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False, ) 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.transpose(1, 2).reshape( batch_size * num_heads, 1, num_queries, num_levels * num_points ) output = ( (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights) .sum(-1) .view(batch_size, num_heads * hidden_dim, num_queries) ) return output.transpose(1, 2).contiguous() class LwDetrMultiscaleDeformableAttention(nn.Module): """ Multiscale deformable attention as proposed in Deformable DETR. """ def __init__(self, config: LwDetrConfig, num_heads: int, n_points: int): super().__init__() self.attn = MultiScaleDeformableAttention() if config.d_model % num_heads != 0: raise ValueError( f"embed_dim (d_model) must be divisible by num_heads, but got {config.d_model} and {num_heads}" ) dim_per_head = config.d_model // num_heads # check if dim_per_head is power of 2 if not ((dim_per_head & (dim_per_head - 1) == 0) and dim_per_head != 0): warnings.warn( "You'd better set embed_dim (d_model) in LwDetrMultiscaleDeformableAttention to make the" " dimension of each attention head a power of 2 which is more efficient in the authors' CUDA" " implementation." ) self.im2col_step = 64 self.d_model = config.d_model self.n_levels = config.num_feature_levels self.n_heads = num_heads self.n_points = n_points self.sampling_offsets = nn.Linear(config.d_model, num_heads * self.n_levels * n_points * 2) self.attention_weights = nn.Linear(config.d_model, num_heads * self.n_levels * n_points) self.value_proj = nn.Linear(config.d_model, config.d_model) self.output_proj = nn.Linear(config.d_model, config.d_model) self.disable_custom_kernels = config.disable_custom_kernels def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, encoder_hidden_states=None, encoder_attention_mask=None, position_embeddings: torch.Tensor | None = None, reference_points=None, spatial_shapes=None, spatial_shapes_list=None, level_start_index=None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor, torch.Tensor]: # add position embeddings to the hidden states before projecting to queries and keys if position_embeddings is not None: hidden_states = hidden_states + position_embeddings batch_size, num_queries, _ = hidden_states.shape batch_size, sequence_length, _ = encoder_hidden_states.shape total_elements = sum(height * width for height, width in spatial_shapes_list) torch_compilable_check( total_elements == sequence_length, "Make sure to align the spatial shapes with the sequence length of the encoder hidden states", ) value = self.value_proj(encoder_hidden_states) if attention_mask is not None: # we invert the attention_mask value = value.masked_fill(~attention_mask[..., None], float(0)) value = value.view(batch_size, sequence_length, self.n_heads, self.d_model // self.n_heads) sampling_offsets = self.sampling_offsets(hidden_states).view( batch_size, num_queries, self.n_heads, self.n_levels, self.n_points, 2 ) attention_weights = self.attention_weights(hidden_states).view( batch_size, num_queries, self.n_heads, self.n_levels * self.n_points ) attention_weights = F.softmax(attention_weights, -1).view( batch_size, num_queries, self.n_heads, self.n_levels, self.n_points ) # batch_size, num_queries, n_heads, n_levels, n_points, 2 num_coordinates = reference_points.shape[-1] if num_coordinates == 2: offset_normalizer = torch.stack([spatial_shapes[..., 1], spatial_shapes[..., 0]], -1) sampling_locations = ( reference_points[:, :, None, :, None, :] + sampling_offsets / offset_normalizer[None, None, None, :, None, :] ) elif num_coordinates == 4: sampling_locations = ( reference_points[:, :, None, :, None, :2] + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5 ) else: raise ValueError(f"Last dim of reference_points must be 2 or 4, but got {reference_points.shape[-1]}") output = self.attn( value, spatial_shapes, spatial_shapes_list, level_start_index, sampling_locations, attention_weights, self.im2col_step, ) output = self.output_proj(output) return output, attention_weights class LwDetrMLP(nn.Module): def __init__(self, config: LwDetrConfig): super().__init__() self.dropout = config.dropout self.activation_fn = ACT2FN[config.decoder_activation_function] self.fc1 = nn.Linear(config.d_model, config.decoder_ffn_dim) self.fc2 = nn.Linear(config.decoder_ffn_dim, config.d_model) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: residual = hidden_states hidden_states = self.fc1(hidden_states) hidden_states = self.activation_fn(hidden_states) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) hidden_states = self.fc2(hidden_states) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) hidden_states = residual + hidden_states return hidden_states class LwDetrDecoderLayer(GradientCheckpointingLayer): def __init__(self, config: LwDetrConfig, layer_idx: int): nn.Module.__init__(self) # self-attention self.self_attn = LwDetrAttention(config, layer_idx=layer_idx) self.dropout = config.dropout self.activation_fn = ACT2FN[config.decoder_activation_function] self.activation_dropout = config.activation_dropout self.self_attn_layer_norm = nn.LayerNorm(config.d_model) # cross-attention self.cross_attn = LwDetrMultiscaleDeformableAttention( config, num_heads=config.decoder_cross_attention_heads, n_points=config.decoder_n_points, ) self.cross_attn_layer_norm = nn.LayerNorm(config.d_model) # mlp self.mlp = LwDetrMLP(config) self.layer_norm = nn.LayerNorm(config.d_model) def forward( self, hidden_states: torch.Tensor, position_embeddings: torch.Tensor | None = None, reference_points=None, spatial_shapes=None, spatial_shapes_list=None, level_start_index=None, encoder_hidden_states: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ): self_attention_output, self_attn_weights = self.self_attn( hidden_states, position_embeddings=position_embeddings, **kwargs ) self_attention_output = nn.functional.dropout(self_attention_output, p=self.dropout, training=self.training) hidden_states = hidden_states + self_attention_output hidden_states = self.self_attn_layer_norm(hidden_states) cross_attention_output, cross_attn_weights = self.cross_attn( hidden_states=hidden_states, attention_mask=encoder_attention_mask, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, position_embeddings=position_embeddings, reference_points=reference_points, spatial_shapes=spatial_shapes, spatial_shapes_list=spatial_shapes_list, level_start_index=level_start_index, **kwargs, ) cross_attention_output = nn.functional.dropout(cross_attention_output, p=self.dropout, training=self.training) hidden_states = hidden_states + cross_attention_output hidden_states = self.cross_attn_layer_norm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = self.layer_norm(hidden_states) return hidden_states @auto_docstring class LwDetrPreTrainedModel(PreTrainedModel): config: LwDetrConfig base_model_prefix = "model" main_input_name = "pixel_values" _no_split_modules = [ r"LwDetrConvEncoder", r"LwDetrDecoderLayer", ] _supports_sdpa = True _supports_flash_attn = True _supports_flex_attn = True _supports_attention_backend = True _can_record_outputs = { "attentions": [LwDetrAttention, LwDetrMultiscaleDeformableAttention], "hidden_states": [LwDetrDecoderLayer], } @torch.no_grad() def _init_weights(self, module): super()._init_weights(module) if isinstance(module, LwDetrMultiscaleDeformableAttention): init.constant_(module.sampling_offsets.weight, 0.0) thetas = torch.arange(module.n_heads, dtype=torch.int64).float() * (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)[0]) .view(module.n_heads, 1, 1, 2) .repeat(1, module.n_levels, module.n_points, 1) ) for i in range(module.n_points): grid_init[:, :, i, :] *= i + 1 init.copy_(module.sampling_offsets.bias, grid_init.view(-1)) init.constant_(module.attention_weights.weight, 0.0) init.constant_(module.attention_weights.bias, 0.0) init.xavier_uniform_(module.value_proj.weight) init.constant_(module.value_proj.bias, 0.0) init.xavier_uniform_(module.output_proj.weight) init.constant_(module.output_proj.bias, 0.0) if hasattr(module, "level_embed"): init.normal_(module.level_embed) if hasattr(module, "refpoint_embed") and module.refpoint_embed is not None: init.constant_(module.refpoint_embed.weight, 0) if hasattr(module, "class_embed") and module.class_embed is not None: prior_prob = 0.01 bias_value = -math.log((1 - prior_prob) / prior_prob) init.constant_(module.class_embed.bias, bias_value) if hasattr(module, "bbox_embed") and module.bbox_embed is not None: init.constant_(module.bbox_embed.layers[-1].weight, 0) init.constant_(module.bbox_embed.layers[-1].bias, 0) @dataclass @auto_docstring( custom_intro=""" Base class for outputs of the LwDetrDecoder. 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 LwDetrDecoderOutput(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. 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_hidden_states: torch.FloatTensor | None = None intermediate_reference_points: torch.FloatTensor | None = None # function to generate sine positional embedding for 4d coordinates def gen_sine_position_embeddings(pos_tensor, hidden_size=256): """ This function computes position embeddings using sine and cosine functions from the input positional tensor, which has a shape of (batch_size, num_queries, 4). The last dimension of `pos_tensor` represents the following coordinates: - 0: x-coord - 1: y-coord - 2: width - 3: height The output shape is (batch_size, num_queries, 512), where final dim (hidden_size*2 = 512) is the total embedding dimension achieved by concatenating the sine and cosine values for each coordinate. """ scale = 2 * math.pi dim = hidden_size // 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) if pos_tensor.size(-1) == 4: w_embed = pos_tensor[:, :, 2] * scale pos_w = w_embed[:, :, None] / dim_t pos_w = torch.stack((pos_w[:, :, 0::2].sin(), pos_w[:, :, 1::2].cos()), dim=3).flatten(2) h_embed = pos_tensor[:, :, 3] * scale pos_h = h_embed[:, :, None] / dim_t pos_h = torch.stack((pos_h[:, :, 0::2].sin(), pos_h[:, :, 1::2].cos()), dim=3).flatten(2) pos = torch.cat((pos_y, pos_x, pos_w, pos_h), dim=2) else: raise ValueError(f"Unknown pos_tensor shape(-1):{pos_tensor.size(-1)}") return pos.to(pos_tensor.dtype) class LwDetrDecoder(LwDetrPreTrainedModel): """ Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`DeformableDetrDecoderLayer`]. The decoder updates the query embeddings through multiple self-attention and deformable cross-attention layers. Some tweaks for LwDetr: - it uses group detr technique at training for faster convergence. Args: config: LwDetrConfig """ def __init__(self, config: LwDetrConfig): super().__init__(config) self.dropout = config.dropout self.layers = nn.ModuleList([LwDetrDecoderLayer(config, i) for i in range(config.decoder_layers)]) self.layernorm = nn.LayerNorm(config.d_model) self.gradient_checkpointing = False self.ref_point_head = LwDetrMLPPredictionHead(2 * config.d_model, config.d_model, config.d_model, num_layers=2) self.post_init() def get_reference(self, reference_points, valid_ratios): # batch_size, num_queries, batch_size, 4 obj_center = reference_points[..., :4] # batch_size, num_queries, num_levels, 4 reference_points_inputs = obj_center[:, :, None] * torch.cat([valid_ratios, valid_ratios], -1)[:, None] # batch_size, num_queries, d_model * 2 query_sine_embed = gen_sine_position_embeddings(reference_points_inputs[:, :, 0, :], self.config.d_model) # batch_size, num_queries, d_model query_pos = self.ref_point_head(query_sine_embed) return reference_points_inputs, query_pos @merge_with_config_defaults @capture_outputs def forward( self, inputs_embeds: torch.Tensor | None = None, reference_points: torch.Tensor | None = None, spatial_shapes: torch.Tensor | None = None, spatial_shapes_list: torch.Tensor | None = None, level_start_index: torch.Tensor | None = None, valid_ratios: torch.Tensor | None = None, encoder_hidden_states: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ): intermediate = () intermediate_reference_points = (reference_points,) if inputs_embeds is not None: hidden_states = inputs_embeds reference_points_inputs, query_pos = self.get_reference(reference_points, valid_ratios) for idx, decoder_layer in enumerate(self.layers): hidden_states = decoder_layer( hidden_states, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, position_embeddings=query_pos, reference_points=reference_points_inputs, spatial_shapes=spatial_shapes, spatial_shapes_list=spatial_shapes_list, level_start_index=level_start_index, **kwargs, ) intermediate_hidden_states = self.layernorm(hidden_states) intermediate += (intermediate_hidden_states,) intermediate = torch.stack(intermediate) last_hidden_state = intermediate[-1] intermediate_reference_points = torch.stack(intermediate_reference_points) return LwDetrDecoderOutput( last_hidden_state=last_hidden_state, intermediate_hidden_states=intermediate, intermediate_reference_points=intermediate_reference_points, ) @dataclass @auto_docstring( custom_intro=""" Base class for outputs of the LwDetr backbone-decoder model. """ ) class LwDetrModelOutput(ModelOutput): r""" init_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`): Initial reference points sent through the Transformer decoder. 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_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). 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. """ init_reference_points: torch.FloatTensor | None = None last_hidden_state: torch.FloatTensor | None = None intermediate_hidden_states: torch.FloatTensor | None = None intermediate_reference_points: torch.FloatTensor | None = None enc_outputs_class: torch.FloatTensor | None = None enc_outputs_coord_logits: torch.FloatTensor | None = None hidden_states: tuple[torch.FloatTensor, ...] | None = None attentions: tuple[torch.FloatTensor, ...] | None = None cross_attentions: tuple[torch.FloatTensor, ...] | None = None def refine_bboxes(reference_points, deltas): reference_points = reference_points.to(deltas.device) new_reference_points_cxcy = deltas[..., :2] * reference_points[..., 2:] + reference_points[..., :2] new_reference_points_wh = deltas[..., 2:].exp() * reference_points[..., 2:] new_reference_points = torch.cat((new_reference_points_cxcy, new_reference_points_wh), -1) return new_reference_points @auto_docstring( custom_intro=""" The bare LW Detr Model (consisting of a backbone and decoder Transformer) outputting raw hidden-states without any specific head on top. """ ) class LwDetrModel(LwDetrPreTrainedModel): def __init__(self, config: LwDetrConfig): super().__init__(config) # Create backbone + positional encoding self.backbone = LwDetrConvEncoder(config) self.group_detr = config.group_detr self.num_queries = config.num_queries hidden_dim = config.d_model self.reference_point_embed = nn.Embedding(self.num_queries * self.group_detr, 4) self.query_feat = nn.Embedding(self.num_queries * self.group_detr, hidden_dim) self.decoder = LwDetrDecoder(config) self.enc_output = nn.ModuleList([nn.Linear(hidden_dim, hidden_dim) for _ in range(self.group_detr)]) self.enc_output_norm = nn.ModuleList([nn.LayerNorm(hidden_dim) for _ in range(self.group_detr)]) # Should normally be None and then instantiated in the ForObjectDetection class self.enc_out_bbox_embed = nn.ModuleList( [LwDetrMLPPredictionHead(config.d_model, config.d_model, 4, num_layers=3) for _ in range(self.group_detr)] ) self.enc_out_class_embed = nn.ModuleList( [nn.Linear(config.d_model, config.num_labels) for _ in range(self.group_detr)] ) self.post_init() def freeze_backbone(self): for name, param in self.backbone.model.named_parameters(): param.requires_grad_(False) def unfreeze_backbone(self): for name, param in self.backbone.model.named_parameters(): param.requires_grad_(True) def get_valid_ratio(self, mask, dtype=torch.float32): """Get the valid ratio of all feature maps.""" _, height, width = mask.shape valid_height = torch.sum(mask[:, :, 0], 1) valid_width = torch.sum(mask[:, 0, :], 1) valid_ratio_height = valid_height.to(dtype) / height valid_ratio_width = valid_width.to(dtype) / width valid_ratio = torch.stack([valid_ratio_width, valid_ratio_height], -1) return valid_ratio def get_proposal_pos_embed(self, proposals): """Get the position embedding of the proposals.""" num_pos_feats = self.config.d_model // 2 temperature = 10000 scale = 2 * math.pi # Compute position embeddings in float32 to avoid overflow with large temperature values in fp16 proposals_dtype = proposals.dtype dim_t = torch.arange(num_pos_feats, dtype=torch.float32, device=proposals.device) dim_t = temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / num_pos_feats) # batch_size, num_queries, 4 proposals = proposals.sigmoid().to(torch.float32) * scale # batch_size, num_queries, 4, 128 pos = proposals[:, :, :, None] / dim_t # batch_size, num_queries, 4, 64, 2 -> batch_size, num_queries, 512 pos = torch.stack((pos[:, :, :, 0::2].sin(), pos[:, :, :, 1::2].cos()), dim=4).flatten(2) # Convert back to target dtype after all computations are done return pos.to(proposals_dtype) def gen_encoder_output_proposals(self, enc_output, padding_mask, spatial_shapes): """Generate the encoder output proposals from encoded enc_output. Args: enc_output (Tensor[batch_size, sequence_length, hidden_size]): Output of the encoder. padding_mask (Tensor[batch_size, sequence_length]): Padding mask for `enc_output`. spatial_shapes (list[tuple[int, int]]): Spatial shapes of the feature maps. Returns: `tuple(torch.FloatTensor)`: A tuple of feature map and bbox prediction. - object_query (Tensor[batch_size, sequence_length, hidden_size]): Object query features. Later used to directly predict a bounding box. (without the need of a decoder) - output_proposals (Tensor[batch_size, sequence_length, 4]): Normalized proposals, after an inverse sigmoid. """ batch_size = enc_output.shape[0] proposals = [] _cur = 0 for level, (height, width) in enumerate(spatial_shapes): mask_flatten_ = padding_mask[:, _cur : (_cur + height * width)].view(batch_size, height, width, 1) valid_height = torch.sum(~mask_flatten_[:, :, 0, 0], 1) valid_width = torch.sum(~mask_flatten_[:, 0, :, 0], 1) grid_y, grid_x = torch.meshgrid( torch.linspace( 0, height - 1, height, dtype=enc_output.dtype, device=enc_output.device, ), torch.linspace( 0, width - 1, width, dtype=enc_output.dtype, device=enc_output.device, ), indexing="ij", ) grid = torch.cat([grid_x.unsqueeze(-1), grid_y.unsqueeze(-1)], -1) scale = torch.cat([valid_width.unsqueeze(-1), valid_height.unsqueeze(-1)], 1).view(batch_size, 1, 1, 2) grid = (grid.unsqueeze(0).expand(batch_size, -1, -1, -1) + 0.5) / scale width_height = torch.ones_like(grid) * 0.05 * (2.0**level) proposal = torch.cat((grid, width_height), -1).view(batch_size, -1, 4) proposals.append(proposal) _cur += height * width output_proposals = torch.cat(proposals, 1) output_proposals_valid = ((output_proposals > 0.01) & (output_proposals < 0.99)).all(-1, keepdim=True) output_proposals = output_proposals.masked_fill(padding_mask.unsqueeze(-1), float("inf")) output_proposals = output_proposals.masked_fill(~output_proposals_valid, float("inf")) # assign each pixel as an object query object_query = enc_output object_query = object_query.masked_fill(padding_mask.unsqueeze(-1), float(0)) object_query = object_query.masked_fill(~output_proposals_valid, float(0)) return object_query, output_proposals @can_return_tuple @auto_docstring def forward( self, pixel_values: torch.FloatTensor = None, pixel_mask: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> LwDetrModelOutput: r""" Examples: ```python >>> from transformers import AutoImageProcessor, DeformableDetrModel >>> from PIL import Image >>> import httpx >>> from io import BytesIO >>> 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("AnnaZhang/lwdetr_small_60e_coco") >>> model = DeformableDetrModel.from_pretrained("AnnaZhang/lwdetr_small_60e_coco") >>> 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] ```""" 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)), dtype=torch.long, device=device) # Extract multi-scale feature maps of same resolution `config.d_model` (cf Figure 4 in paper) # First, sent pixel_values + pixel_mask through Backbone to obtain the features # which is a list of tuples features = self.backbone(pixel_values, pixel_mask) # Then, apply 1x1 convolution to reduce the channel dimension to d_model (256 by default) sources = [] masks = [] for level, (source, mask) in enumerate(features): sources.append(source) masks.append(mask) if mask is None: raise ValueError("No attention mask was provided") if self.training: reference_points = self.reference_point_embed.weight query_feat = self.query_feat.weight else: # only use one group in inference reference_points = self.reference_point_embed.weight[: self.num_queries] query_feat = self.query_feat.weight[: self.num_queries] # Prepare encoder inputs (by flattening) source_flatten = [] mask_flatten = [] spatial_shapes_list = [] for source, mask in zip(sources, masks): batch_size, num_channels, height, width = source.shape spatial_shape = (height, width) spatial_shapes_list.append(spatial_shape) source = source.flatten(2).transpose(1, 2) mask = mask.flatten(1) source_flatten.append(source) mask_flatten.append(mask) source_flatten = torch.cat(source_flatten, 1) mask_flatten = torch.cat(mask_flatten, 1) spatial_shapes = torch.as_tensor(spatial_shapes_list, dtype=torch.long, device=source_flatten.device) level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1])) valid_ratios = torch.stack([self.get_valid_ratio(m, dtype=source_flatten.dtype) for m in masks], 1) target = query_feat.unsqueeze(0).expand(batch_size, -1, -1) reference_points = reference_points.unsqueeze(0).expand(batch_size, -1, -1) object_query_embedding, output_proposals = self.gen_encoder_output_proposals( source_flatten, ~mask_flatten, spatial_shapes_list ) group_detr = self.group_detr if self.training else 1 topk = self.num_queries topk_coords_logits = [] topk_coords_logits_undetach = [] object_query_undetach = [] for group_id in range(group_detr): group_object_query = self.enc_output[group_id](object_query_embedding) group_object_query = self.enc_output_norm[group_id](group_object_query) group_enc_outputs_class = self.enc_out_class_embed[group_id](group_object_query) group_delta_bbox = self.enc_out_bbox_embed[group_id](group_object_query) group_enc_outputs_coord = refine_bboxes(output_proposals, group_delta_bbox) group_topk_proposals = torch.topk(group_enc_outputs_class.max(-1)[0], topk, dim=1)[1] group_topk_coords_logits_undetach = torch.gather( group_enc_outputs_coord, 1, group_topk_proposals.unsqueeze(-1).repeat(1, 1, 4), ) group_topk_coords_logits = group_topk_coords_logits_undetach.detach() group_object_query_undetach = torch.gather( group_object_query, 1, group_topk_proposals.unsqueeze(-1).repeat(1, 1, self.config.d_model) ) topk_coords_logits.append(group_topk_coords_logits) topk_coords_logits_undetach.append(group_topk_coords_logits_undetach) object_query_undetach.append(group_object_query_undetach) topk_coords_logits = torch.cat(topk_coords_logits, 1) topk_coords_logits_undetach = torch.cat(topk_coords_logits_undetach, 1) object_query_undetach = torch.cat(object_query_undetach, 1) enc_outputs_class = object_query_undetach enc_outputs_coord_logits = topk_coords_logits reference_points = refine_bboxes(topk_coords_logits_undetach, reference_points) init_reference_points = reference_points decoder_outputs = self.decoder( inputs_embeds=target, reference_points=reference_points, spatial_shapes=spatial_shapes, spatial_shapes_list=spatial_shapes_list, level_start_index=level_start_index, valid_ratios=valid_ratios, encoder_hidden_states=source_flatten, encoder_attention_mask=mask_flatten, **kwargs, ) return LwDetrModelOutput( init_reference_points=init_reference_points, last_hidden_state=decoder_outputs.last_hidden_state, intermediate_hidden_states=decoder_outputs.intermediate_hidden_states, intermediate_reference_points=decoder_outputs.intermediate_reference_points, enc_outputs_class=enc_outputs_class, enc_outputs_coord_logits=enc_outputs_coord_logits, hidden_states=decoder_outputs.hidden_states, attentions=decoder_outputs.attentions, cross_attentions=decoder_outputs.cross_attentions, ) class LwDetrMLPPredictionHead(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 @dataclass @auto_docstring( custom_intro=""" Output type of [`LwDetrForObjectDetection`]. """ ) class LwDetrObjectDetectionOutput(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 [`~DeformableDetrProcessor.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. init_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`): Initial reference points sent through the Transformer decoder. 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_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). 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. """ 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 init_reference_points: torch.FloatTensor | None = None last_hidden_state: torch.FloatTensor | None = None intermediate_hidden_states: torch.FloatTensor | None = None intermediate_reference_points: torch.FloatTensor | None = None enc_outputs_class: Any = None enc_outputs_coord_logits: torch.FloatTensor | None = None hidden_states: tuple[torch.FloatTensor, ...] | None = None attentions: tuple[torch.FloatTensor, ...] | None = None cross_attentions: tuple[torch.FloatTensor, ...] | None = None @auto_docstring( custom_intro=""" LW DETR Model (consisting of a backbone and decoder Transformer) with object detection heads on top, for tasks such as COCO detection. """ ) class LwDetrForObjectDetection(LwDetrPreTrainedModel): # When using clones, all layers > 0 will be clones, but layer 0 *is* required # We can't initialize the model on meta device as some weights are modified during the initialization _no_split_modules = None _tied_weights_keys = None def __init__(self, config: LwDetrConfig): super().__init__(config) self.model = LwDetrModel(config) self.class_embed = nn.Linear(config.d_model, config.num_labels) self.bbox_embed = LwDetrMLPPredictionHead(config.d_model, config.d_model, 4, num_layers=3) self.post_init() @can_return_tuple @auto_docstring def forward( self, pixel_values: torch.FloatTensor = None, pixel_mask: torch.LongTensor | None = None, labels: list[dict] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> LwDetrObjectDetectionOutput: 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, LwDetrForObjectDetection >>> from PIL import Image >>> import httpx >>> from io import BytesIO >>> 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("AnnaZhang/lwdetr_small_60e_coco") >>> model = LwDetrForObjectDetection.from_pretrained("AnnaZhang/lwdetr_small_60e_coco") >>> 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 cat with confidence 0.8 at location [16.5, 52.84, 318.25, 470.78] Detected cat with confidence 0.789 at location [342.19, 24.3, 640.02, 372.25] Detected remote with confidence 0.633 at location [40.79, 72.78, 176.76, 117.25] ```""" outputs = self.model( pixel_values, pixel_mask=pixel_mask, **kwargs, ) last_hidden_states = outputs.last_hidden_state intermediate_reference_points = outputs.intermediate_reference_points enc_outputs_class_logits = outputs.enc_outputs_class enc_outputs_boxes_logits = outputs.enc_outputs_coord_logits logits = self.class_embed(last_hidden_states) pred_boxes_delta = self.bbox_embed(last_hidden_states) pred_boxes = refine_bboxes(intermediate_reference_points[-1], pred_boxes_delta) enc_outputs_class_logits_list = enc_outputs_class_logits.split(self.config.num_queries, dim=1) pred_class = [] group_detr = self.config.group_detr if self.training else 1 for group_index in range(group_detr): group_pred_class = self.model.enc_out_class_embed[group_index](enc_outputs_class_logits_list[group_index]) pred_class.append(group_pred_class) enc_outputs_class_logits = torch.cat(pred_class, dim=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_hidden_states = outputs.intermediate_hidden_states outputs_coord_delta = self.bbox_embed(intermediate_hidden_states) outputs_coord = refine_bboxes(intermediate_reference_points, outputs_coord_delta) outputs_class = self.class_embed(intermediate_hidden_states) loss, loss_dict, auxiliary_outputs = self.loss_function( logits, labels, self.device, pred_boxes, self.config, outputs_class, outputs_coord, enc_outputs_class_logits, enc_outputs_boxes_logits, ) return LwDetrObjectDetectionOutput( 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_reference_points=outputs.intermediate_reference_points, init_reference_points=outputs.init_reference_points, enc_outputs_class=enc_outputs_class_logits, enc_outputs_coord_logits=enc_outputs_boxes_logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions, cross_attentions=outputs.cross_attentions, ) __all__ = [ "LwDetrPreTrainedModel", "LwDetrModel", "LwDetrForObjectDetection", "LwDetrViTPreTrainedModel", "LwDetrViTBackbone", ]