dgenerate-ultralytics-headless 8.3.196__py3-none-any.whl → 8.3.248__py3-none-any.whl

ultralytics/nn/modules/transformer.py

@@ -10,26 +10,27 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch.nn.init import constant_, xavier_uniform_
 
+from ultralytics.utils.torch_utils import TORCH_1_11
+
 from .conv import Conv
 from .utils import _get_clones, inverse_sigmoid, multi_scale_deformable_attn_pytorch
 
 __all__ = (
-    "TransformerEncoderLayer",
-    "TransformerLayer",
-    "TransformerBlock",
-    "MLPBlock",
-    "LayerNorm2d",
     "AIFI",
+    "MLP",
     "DeformableTransformerDecoder",
     "DeformableTransformerDecoderLayer",
+    "LayerNorm2d",
+    "MLPBlock",
     "MSDeformAttn",
-    "MLP",
+    "TransformerBlock",
+    "TransformerEncoderLayer",
+    "TransformerLayer",
 )
 
 
 class TransformerEncoderLayer(nn.Module):
-    """
-    A single layer of the transformer encoder.
+    """A single layer of the transformer encoder.
 
     This class implements a standard transformer encoder layer with multi-head attention and feedforward network,
     supporting both pre-normalization and post-normalization configurations.
@@ -56,8 +57,7 @@ class TransformerEncoderLayer(nn.Module):
         act: nn.Module = nn.GELU(),
         normalize_before: bool = False,
     ):
-        """
-        Initialize the TransformerEncoderLayer with specified parameters.
+        """Initialize the TransformerEncoderLayer with specified parameters.
 
         Args:
             c1 (int): Input dimension.
@@ -100,8 +100,7 @@ class TransformerEncoderLayer(nn.Module):
         src_key_padding_mask: torch.Tensor | None = None,
         pos: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """
-        Perform forward pass with post-normalization.
+        """Perform forward pass with post-normalization.
 
         Args:
             src (torch.Tensor): Input tensor.
@@ -127,8 +126,7 @@ class TransformerEncoderLayer(nn.Module):
         src_key_padding_mask: torch.Tensor | None = None,
         pos: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """
-        Perform forward pass with pre-normalization.
+        """Perform forward pass with pre-normalization.
 
         Args:
             src (torch.Tensor): Input tensor.
@@ -154,8 +152,7 @@ class TransformerEncoderLayer(nn.Module):
         src_key_padding_mask: torch.Tensor | None = None,
         pos: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """
-        Forward propagate the input through the encoder module.
+        """Forward propagate the input through the encoder module.
 
         Args:
             src (torch.Tensor): Input tensor.
@@ -172,8 +169,7 @@ class TransformerEncoderLayer(nn.Module):
 
 
 class AIFI(TransformerEncoderLayer):
-    """
-    AIFI transformer layer for 2D data with positional embeddings.
+    """AIFI transformer layer for 2D data with positional embeddings.
 
     This class extends TransformerEncoderLayer to work with 2D feature maps by adding 2D sine-cosine positional
     embeddings and handling the spatial dimensions appropriately.
@@ -188,8 +184,7 @@ class AIFI(TransformerEncoderLayer):
         act: nn.Module = nn.GELU(),
         normalize_before: bool = False,
     ):
-        """
-        Initialize the AIFI instance with specified parameters.
+        """Initialize the AIFI instance with specified parameters.
 
         Args:
             c1 (int): Input dimension.
@@ -202,8 +197,7 @@ class AIFI(TransformerEncoderLayer):
         super().__init__(c1, cm, num_heads, dropout, act, normalize_before)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Forward pass for the AIFI transformer layer.
+        """Forward pass for the AIFI transformer layer.
 
         Args:
             x (torch.Tensor): Input tensor with shape [B, C, H, W].
@@ -221,8 +215,7 @@ class AIFI(TransformerEncoderLayer):
     def build_2d_sincos_position_embedding(
         w: int, h: int, embed_dim: int = 256, temperature: float = 10000.0
     ) -> torch.Tensor:
-        """
-        Build 2D sine-cosine position embedding.
+        """Build 2D sine-cosine position embedding.
 
         Args:
             w (int): Width of the feature map.
@@ -236,7 +229,7 @@ class AIFI(TransformerEncoderLayer):
         assert embed_dim % 4 == 0, "Embed dimension must be divisible by 4 for 2D sin-cos position embedding"
         grid_w = torch.arange(w, dtype=torch.float32)
         grid_h = torch.arange(h, dtype=torch.float32)
-        grid_w, grid_h = torch.meshgrid(grid_w, grid_h, indexing="ij")
+        grid_w, grid_h = torch.meshgrid(grid_w, grid_h, indexing="ij") if TORCH_1_11 else torch.meshgrid(grid_w, grid_h)
         pos_dim = embed_dim // 4
         omega = torch.arange(pos_dim, dtype=torch.float32) / pos_dim
         omega = 1.0 / (temperature**omega)
@@ -251,8 +244,7 @@ class TransformerLayer(nn.Module):
     """Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance)."""
 
     def __init__(self, c: int, num_heads: int):
-        """
-        Initialize a self-attention mechanism using linear transformations and multi-head attention.
+        """Initialize a self-attention mechanism using linear transformations and multi-head attention.
 
         Args:
             c (int): Input and output channel dimension.
@@ -267,8 +259,7 @@ class TransformerLayer(nn.Module):
         self.fc2 = nn.Linear(c, c, bias=False)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Apply a transformer block to the input x and return the output.
+        """Apply a transformer block to the input x and return the output.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -281,11 +272,10 @@ class TransformerLayer(nn.Module):
 
 
 class TransformerBlock(nn.Module):
-    """
-    Vision Transformer block based on https://arxiv.org/abs/2010.11929.
+    """Vision Transformer block based on https://arxiv.org/abs/2010.11929.
 
-    This class implements a complete transformer block with optional convolution layer for channel adjustment,
-    learnable position embedding, and multiple transformer layers.
+    This class implements a complete transformer block with optional convolution layer for channel adjustment, learnable
+    position embedding, and multiple transformer layers.
 
     Attributes:
         conv (Conv, optional): Convolution layer if input and output channels differ.
@@ -295,8 +285,7 @@ class TransformerBlock(nn.Module):
     """
 
     def __init__(self, c1: int, c2: int, num_heads: int, num_layers: int):
-        """
-        Initialize a Transformer module with position embedding and specified number of heads and layers.
+        """Initialize a Transformer module with position embedding and specified number of heads and layers.
 
         Args:
             c1 (int): Input channel dimension.
@@ -313,28 +302,26 @@ class TransformerBlock(nn.Module):
         self.c2 = c2
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Forward propagate the input through the transformer block.
+        """Forward propagate the input through the transformer block.
 
         Args:
-            x (torch.Tensor): Input tensor with shape [b, c1, w, h].
+            x (torch.Tensor): Input tensor with shape [b, c1, h, w].
 
         Returns:
-            (torch.Tensor): Output tensor with shape [b, c2, w, h].
+            (torch.Tensor): Output tensor with shape [b, c2, h, w].
         """
         if self.conv is not None:
             x = self.conv(x)
-        b, _, w, h = x.shape
+        b, _, h, w = x.shape
         p = x.flatten(2).permute(2, 0, 1)
-        return self.tr(p + self.linear(p)).permute(1, 2, 0).reshape(b, self.c2, w, h)
+        return self.tr(p + self.linear(p)).permute(1, 2, 0).reshape(b, self.c2, h, w)
 
 
 class MLPBlock(nn.Module):
     """A single block of a multi-layer perceptron."""
 
     def __init__(self, embedding_dim: int, mlp_dim: int, act=nn.GELU):
-        """
-        Initialize the MLPBlock with specified embedding dimension, MLP dimension, and activation function.
+        """Initialize the MLPBlock with specified embedding dimension, MLP dimension, and activation function.
 
         Args:
             embedding_dim (int): Input and output dimension.
@@ -347,8 +334,7 @@ class MLPBlock(nn.Module):
         self.act = act()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Forward pass for the MLPBlock.
+        """Forward pass for the MLPBlock.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -360,11 +346,10 @@ class MLPBlock(nn.Module):
 
 
 class MLP(nn.Module):
-    """
-    A simple multi-layer perceptron (also called FFN).
+    """A simple multi-layer perceptron (also called FFN).
 
-    This class implements a configurable MLP with multiple linear layers, activation functions, and optional
-    sigmoid output activation.
+    This class implements a configurable MLP with multiple linear layers, activation functions, and optional sigmoid
+    output activation.
 
     Attributes:
         num_layers (int): Number of layers in the MLP.
@@ -374,10 +359,17 @@ class MLP(nn.Module):
     """
 
     def __init__(
-        self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int, act=nn.ReLU, sigmoid: bool = False
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        output_dim: int,
+        num_layers: int,
+        act=nn.ReLU,
+        sigmoid: bool = False,
+        residual: bool = False,
+        out_norm: nn.Module = None,
     ):
-        """
-        Initialize the MLP with specified input, hidden, output dimensions and number of layers.
+        """Initialize the MLP with specified input, hidden, output dimensions and number of layers.
 
         Args:
             input_dim (int): Input dimension.
@@ -386,17 +378,24 @@ class MLP(nn.Module):
             num_layers (int): Number of layers.
             act (nn.Module): Activation function.
             sigmoid (bool): Whether to apply sigmoid to the output.
+            residual (bool): Whether to use residual connections.
+            out_norm (nn.Module, optional): Normalization layer for the output.
         """
         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]))
+        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim, *h], [*h, output_dim]))
         self.sigmoid = sigmoid
         self.act = act()
+        if residual and input_dim != output_dim:
+            raise ValueError("residual is only supported if input_dim == output_dim")
+        self.residual = residual
+        # whether to apply a normalization layer to the output
+        assert isinstance(out_norm, nn.Module) or out_norm is None
+        self.out_norm = out_norm or nn.Identity()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Forward pass for the entire MLP.
+        """Forward pass for the entire MLP.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -404,17 +403,20 @@ class MLP(nn.Module):
         Returns:
             (torch.Tensor): Output tensor after MLP.
         """
+        orig_x = x
         for i, layer in enumerate(self.layers):
             x = getattr(self, "act", nn.ReLU())(layer(x)) if i < self.num_layers - 1 else layer(x)
+        if getattr(self, "residual", False):
+            x = x + orig_x
+        x = getattr(self, "out_norm", nn.Identity())(x)
         return x.sigmoid() if getattr(self, "sigmoid", False) else x
 
 
 class LayerNorm2d(nn.Module):
-    """
-    2D Layer Normalization module inspired by Detectron2 and ConvNeXt implementations.
+    """2D Layer Normalization module inspired by Detectron2 and ConvNeXt implementations.
 
-    This class implements layer normalization for 2D feature maps, normalizing across the channel dimension
-    while preserving spatial dimensions.
+    This class implements layer normalization for 2D feature maps, normalizing across the channel dimension while
+    preserving spatial dimensions.
 
     Attributes:
         weight (nn.Parameter): Learnable scale parameter.
@@ -427,8 +429,7 @@ class LayerNorm2d(nn.Module):
     """
 
     def __init__(self, num_channels: int, eps: float = 1e-6):
-        """
-        Initialize LayerNorm2d with the given parameters.
+        """Initialize LayerNorm2d with the given parameters.
 
         Args:
             num_channels (int): Number of channels in the input.
@@ -440,8 +441,7 @@ class LayerNorm2d(nn.Module):
         self.eps = eps
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Perform forward pass for 2D layer normalization.
+        """Perform forward pass for 2D layer normalization.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -456,11 +456,10 @@ class LayerNorm2d(nn.Module):
 
 
 class MSDeformAttn(nn.Module):
-    """
-    Multiscale Deformable Attention Module based on Deformable-DETR and PaddleDetection implementations.
+    """Multiscale Deformable Attention Module based on Deformable-DETR and PaddleDetection implementations.
 
-    This module implements multiscale deformable attention that can attend to features at multiple scales
-    with learnable sampling locations and attention weights.
+    This module implements multiscale deformable attention that can attend to features at multiple scales with learnable
+    sampling locations and attention weights.
 
     Attributes:
         im2col_step (int): Step size for im2col operations.
@@ -478,8 +477,7 @@ class MSDeformAttn(nn.Module):
     """
 
     def __init__(self, d_model: int = 256, n_levels: int = 4, n_heads: int = 8, n_points: int = 4):
-        """
-        Initialize MSDeformAttn with the given parameters.
+        """Initialize MSDeformAttn with the given parameters.
 
         Args:
             d_model (int): Model dimension.
@@ -537,13 +535,12 @@ class MSDeformAttn(nn.Module):
         value_shapes: list,
         value_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """
-        Perform forward pass for multiscale deformable attention.
+        """Perform forward pass for multiscale deformable attention.
 
         Args:
             query (torch.Tensor): Query tensor with shape [bs, query_length, C].
-            refer_bbox (torch.Tensor): Reference bounding boxes with shape [bs, query_length, n_levels, 2],
-                range in [0, 1], top-left (0,0), bottom-right (1, 1), including padding area.
+            refer_bbox (torch.Tensor): Reference bounding boxes with shape [bs, query_length, n_levels, 2], range in [0,
+                1], top-left (0,0), bottom-right (1, 1), including padding area.
             value (torch.Tensor): Value tensor with shape [bs, value_length, C].
             value_shapes (list): List with shape [n_levels, 2], [(H_0, W_0), (H_1, W_1), ..., (H_{L-1}, W_{L-1})].
             value_mask (torch.Tensor, optional): Mask tensor with shape [bs, value_length], True for non-padding
@@ -582,8 +579,7 @@ class MSDeformAttn(nn.Module):
 
 
 class DeformableTransformerDecoderLayer(nn.Module):
-    """
-    Deformable Transformer Decoder Layer inspired by PaddleDetection and Deformable-DETR implementations.
+    """Deformable Transformer Decoder Layer inspired by PaddleDetection and Deformable-DETR implementations.
 
     This class implements a single decoder layer with self-attention, cross-attention using multiscale deformable
     attention, and a feedforward network.
@@ -617,8 +613,7 @@ class DeformableTransformerDecoderLayer(nn.Module):
         n_levels: int = 4,
         n_points: int = 4,
     ):
-        """
-        Initialize the DeformableTransformerDecoderLayer with the given parameters.
+        """Initialize the DeformableTransformerDecoderLayer with the given parameters.
 
         Args:
             d_model (int): Model dimension.
@@ -655,8 +650,7 @@ class DeformableTransformerDecoderLayer(nn.Module):
         return tensor if pos is None else tensor + pos
 
     def forward_ffn(self, tgt: torch.Tensor) -> torch.Tensor:
-        """
-        Perform forward pass through the Feed-Forward Network part of the layer.
+        """Perform forward pass through the Feed-Forward Network part of the layer.
 
         Args:
             tgt (torch.Tensor): Input tensor.
@@ -678,8 +672,7 @@ class DeformableTransformerDecoderLayer(nn.Module):
         attn_mask: torch.Tensor | None = None,
         query_pos: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """
-        Perform the forward pass through the entire decoder layer.
+        """Perform the forward pass through the entire decoder layer.
 
         Args:
             embed (torch.Tensor): Input embeddings.
@@ -713,11 +706,10 @@ class DeformableTransformerDecoderLayer(nn.Module):
 
 
 class DeformableTransformerDecoder(nn.Module):
-    """
-    Deformable Transformer Decoder based on PaddleDetection implementation.
+    """Deformable Transformer Decoder based on PaddleDetection implementation.
 
-    This class implements a complete deformable transformer decoder with multiple decoder layers and prediction
-    heads for bounding box regression and classification.
+    This class implements a complete deformable transformer decoder with multiple decoder layers and prediction heads
+    for bounding box regression and classification.
 
     Attributes:
         layers (nn.ModuleList): List of decoder layers.
@@ -730,8 +722,7 @@ class DeformableTransformerDecoder(nn.Module):
     """
 
     def __init__(self, hidden_dim: int, decoder_layer: nn.Module, num_layers: int, eval_idx: int = -1):
-        """
-        Initialize the DeformableTransformerDecoder with the given parameters.
+        """Initialize the DeformableTransformerDecoder with the given parameters.
 
         Args:
             hidden_dim (int): Hidden dimension.
@@ -757,8 +748,7 @@ class DeformableTransformerDecoder(nn.Module):
         attn_mask: torch.Tensor | None = None,
         padding_mask: torch.Tensor | None = None,
     ):
-        """
-        Perform the forward pass through the entire decoder.
+        """Perform the forward pass through the entire decoder.
 
         Args:
             embed (torch.Tensor): Decoder embeddings.

ultralytics/nn/modules/utils.py

@@ -9,12 +9,11 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch.nn.init import uniform_
 
-__all__ = "multi_scale_deformable_attn_pytorch", "inverse_sigmoid"
+__all__ = "inverse_sigmoid", "multi_scale_deformable_attn_pytorch"
 
 
 def _get_clones(module, n):
-    """
-    Create a list of cloned modules from the given module.
+    """Create a list of cloned modules from the given module.
 
     Args:
         module (nn.Module): The module to be cloned.
@@ -34,12 +33,11 @@ def _get_clones(module, n):
 
 
 def bias_init_with_prob(prior_prob=0.01):
-    """
-    Initialize conv/fc bias value according to a given probability value.
+    """Initialize conv/fc bias value according to a given probability value.
 
-    This function calculates the bias initialization value based on a prior probability using the inverse error function.
-    It's commonly used in object detection models to initialize classification layers with a specific positive prediction
-    probability.
+    This function calculates the bias initialization value based on a prior probability using the inverse error
+    function. It's commonly used in object detection models to initialize classification layers with a specific positive
+    prediction probability.
 
     Args:
         prior_prob (float, optional): Prior probability for bias initialization.
@@ -56,11 +54,10 @@ def bias_init_with_prob(prior_prob=0.01):
 
 
 def linear_init(module):
-    """
-    Initialize the weights and biases of a linear module.
+    """Initialize the weights and biases of a linear module.
 
-    This function initializes the weights of a linear module using a uniform distribution within bounds calculated
-    from the input dimension. If the module has a bias, it is also initialized.
+    This function initializes the weights of a linear module using a uniform distribution within bounds calculated from
+    the input dimension. If the module has a bias, it is also initialized.
 
     Args:
         module (nn.Module): Linear module to initialize.
@@ -80,8 +77,7 @@ def linear_init(module):
 
 
 def inverse_sigmoid(x, eps=1e-5):
-    """
-    Calculate the inverse sigmoid function for a tensor.
+    """Calculate the inverse sigmoid function for a tensor.
 
     This function applies the inverse of the sigmoid function to a tensor, which is useful in various neural network
     operations, particularly in attention mechanisms and coordinate transformations.
@@ -110,8 +106,7 @@ def multi_scale_deformable_attn_pytorch(
     sampling_locations: torch.Tensor,
     attention_weights: torch.Tensor,
 ) -> torch.Tensor:
-    """
-    Implement multi-scale deformable attention in PyTorch.
+    """Implement multi-scale deformable attention in PyTorch.
 
     This function performs deformable attention across multiple feature map scales, allowing the model to attend to
     different spatial locations with learned offsets.
@@ -119,10 +114,10 @@ def multi_scale_deformable_attn_pytorch(
     Args:
         value (torch.Tensor): The value tensor with shape (bs, num_keys, num_heads, embed_dims).
         value_spatial_shapes (torch.Tensor): Spatial shapes of the value tensor with shape (num_levels, 2).
-        sampling_locations (torch.Tensor): The sampling locations with shape
-            (bs, num_queries, num_heads, num_levels, num_points, 2).
-        attention_weights (torch.Tensor): The attention weights with shape
-            (bs, num_queries, num_heads, num_levels, num_points).
+        sampling_locations (torch.Tensor): The sampling locations with shape (bs, num_queries, num_heads, num_levels,
+            num_points, 2).
+        attention_weights (torch.Tensor): The attention weights with shape (bs, num_queries, num_heads, num_levels,
+            num_points).
 
     Returns:
         (torch.Tensor): The output tensor with shape (bs, num_queries, embed_dims).
@@ -152,7 +147,7 @@ def multi_scale_deformable_attn_pytorch(
         sampling_value_list.append(sampling_value_l_)
     # (bs, num_queries, num_heads, num_levels, num_points) ->
     # (bs, num_heads, num_queries, num_levels, num_points) ->
-    # (bs, num_heads, 1, num_queries, num_levels*num_points)
+    # (bs*num_heads, 1, num_queries, num_levels*num_points)
     attention_weights = attention_weights.transpose(1, 2).reshape(
         bs * num_heads, 1, num_queries, num_levels * num_points
     )