ultralytics 8.3.143__py3-none-any.whl → 8.3.144__py3-none-any.whl

ultralytics/nn/modules/transformer.py

@@ -2,6 +2,7 @@
 """Transformer modules."""
 
 import math
+from typing import List, Optional
 
 import torch
 import torch.nn as nn
@@ -27,7 +28,10 @@ __all__ = (
 
 class TransformerEncoderLayer(nn.Module):
     """
-    Defines 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.
 
     Attributes:
         ma (nn.MultiheadAttention): Multi-head attention module.
@@ -42,7 +46,15 @@ class TransformerEncoderLayer(nn.Module):
         normalize_before (bool): Whether to apply normalization before attention and feedforward.
     """
 
-    def __init__(self, c1, cm=2048, num_heads=8, dropout=0.0, act=nn.GELU(), normalize_before=False):
+    def __init__(
+        self,
+        c1: int,
+        cm: int = 2048,
+        num_heads: int = 8,
+        dropout: float = 0.0,
+        act: nn.Module = nn.GELU(),
+        normalize_before: bool = False,
+    ):
         """
         Initialize the TransformerEncoderLayer with specified parameters.
 
@@ -76,11 +88,17 @@ class TransformerEncoderLayer(nn.Module):
         self.normalize_before = normalize_before
 
     @staticmethod
-    def with_pos_embed(tensor, pos=None):
+    def with_pos_embed(tensor: torch.Tensor, pos: Optional[torch.Tensor] = None) -> torch.Tensor:
         """Add position embeddings to the tensor if provided."""
         return tensor if pos is None else tensor + pos
 
-    def forward_post(self, src, src_mask=None, src_key_padding_mask=None, pos=None):
+    def forward_post(
+        self,
+        src: torch.Tensor,
+        src_mask: Optional[torch.Tensor] = None,
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+        pos: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
         """
         Perform forward pass with post-normalization.
 
@@ -101,7 +119,13 @@ class TransformerEncoderLayer(nn.Module):
         src = src + self.dropout2(src2)
         return self.norm2(src)
 
-    def forward_pre(self, src, src_mask=None, src_key_padding_mask=None, pos=None):
+    def forward_pre(
+        self,
+        src: torch.Tensor,
+        src_mask: Optional[torch.Tensor] = None,
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+        pos: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
         """
         Perform forward pass with pre-normalization.
 
@@ -122,9 +146,15 @@ class TransformerEncoderLayer(nn.Module):
         src2 = self.fc2(self.dropout(self.act(self.fc1(src2))))
         return src + self.dropout2(src2)
 
-    def forward(self, src, src_mask=None, src_key_padding_mask=None, pos=None):
+    def forward(
+        self,
+        src: torch.Tensor,
+        src_mask: Optional[torch.Tensor] = None,
+        src_key_padding_mask: Optional[torch.Tensor] = None,
+        pos: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
         """
-        Forward propagates the input through the encoder module.
+        Forward propagate the input through the encoder module.
 
         Args:
             src (torch.Tensor): Input tensor.
@@ -142,12 +172,21 @@ class TransformerEncoderLayer(nn.Module):
 
 class AIFI(TransformerEncoderLayer):
     """
-    Defines the AIFI transformer layer.
+    AIFI transformer layer for 2D data with positional embeddings.
 
-    This class extends TransformerEncoderLayer to work with 2D data by adding 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.
     """
 
-    def __init__(self, c1, cm=2048, num_heads=8, dropout=0, act=nn.GELU(), normalize_before=False):
+    def __init__(
+        self,
+        c1: int,
+        cm: int = 2048,
+        num_heads: int = 8,
+        dropout: float = 0,
+        act: nn.Module = nn.GELU(),
+        normalize_before: bool = False,
+    ):
         """
         Initialize the AIFI instance with specified parameters.
 
@@ -161,7 +200,7 @@ class AIFI(TransformerEncoderLayer):
         """
         super().__init__(c1, cm, num_heads, dropout, act, normalize_before)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Forward pass for the AIFI transformer layer.
 
@@ -178,7 +217,9 @@ class AIFI(TransformerEncoderLayer):
         return x.permute(0, 2, 1).view([-1, c, h, w]).contiguous()
 
     @staticmethod
-    def build_2d_sincos_position_embedding(w, h, embed_dim=256, temperature=10000.0):
+    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.
 
@@ -208,7 +249,7 @@ class AIFI(TransformerEncoderLayer):
 class TransformerLayer(nn.Module):
     """Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance)."""
 
-    def __init__(self, c, num_heads):
+    def __init__(self, c: int, num_heads: int):
         """
         Initialize a self-attention mechanism using linear transformations and multi-head attention.
 
@@ -224,7 +265,7 @@ class TransformerLayer(nn.Module):
         self.fc1 = nn.Linear(c, c, bias=False)
         self.fc2 = nn.Linear(c, c, bias=False)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Apply a transformer block to the input x and return the output.
 
@@ -240,7 +281,10 @@ class TransformerLayer(nn.Module):
 
 class TransformerBlock(nn.Module):
     """
-    Vision Transformer 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.
 
     Attributes:
         conv (Conv, optional): Convolution layer if input and output channels differ.
@@ -249,7 +293,7 @@ class TransformerBlock(nn.Module):
         c2 (int): Output channel dimension.
     """
 
-    def __init__(self, c1, c2, num_heads, num_layers):
+    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.
 
@@ -267,9 +311,9 @@ class TransformerBlock(nn.Module):
         self.tr = nn.Sequential(*(TransformerLayer(c2, num_heads) for _ in range(num_layers)))
         self.c2 = c2
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
-        Forward propagates the input through the bottleneck module.
+        Forward propagate the input through the transformer block.
 
         Args:
             x (torch.Tensor): Input tensor with shape [b, c1, w, h].
@@ -285,9 +329,9 @@ class TransformerBlock(nn.Module):
 
 
 class MLPBlock(nn.Module):
-    """Implements a single block of a multi-layer perceptron."""
+    """A single block of a multi-layer perceptron."""
 
-    def __init__(self, embedding_dim, mlp_dim, act=nn.GELU):
+    def __init__(self, embedding_dim: int, mlp_dim: int, act=nn.GELU):
         """
         Initialize the MLPBlock with specified embedding dimension, MLP dimension, and activation function.
 
@@ -316,7 +360,10 @@ class MLPBlock(nn.Module):
 
 class MLP(nn.Module):
     """
-    Implements 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.
 
     Attributes:
         num_layers (int): Number of layers in the MLP.
@@ -325,7 +372,9 @@ class MLP(nn.Module):
         act (nn.Module): Activation function.
     """
 
-    def __init__(self, input_dim, hidden_dim, output_dim, num_layers, act=nn.ReLU, sigmoid=False):
+    def __init__(
+        self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int, act=nn.ReLU, sigmoid: bool = False
+    ):
         """
         Initialize the MLP with specified input, hidden, output dimensions and number of layers.
 
@@ -344,7 +393,7 @@ class MLP(nn.Module):
         self.sigmoid = sigmoid
         self.act = act()
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Forward pass for the entire MLP.
 
@@ -363,18 +412,20 @@ class LayerNorm2d(nn.Module):
     """
     2D Layer Normalization module inspired by Detectron2 and ConvNeXt implementations.
 
-    Original implementations in
-    https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py
-    and
-    https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext.py.
+    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.
         bias (nn.Parameter): Learnable bias parameter.
         eps (float): Small constant for numerical stability.
+
+    References:
+        https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py
+        https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext.py
     """
 
-    def __init__(self, num_channels, eps=1e-6):
+    def __init__(self, num_channels: int, eps: float = 1e-6):
         """
         Initialize LayerNorm2d with the given parameters.
 
@@ -387,7 +438,7 @@ class LayerNorm2d(nn.Module):
         self.bias = nn.Parameter(torch.zeros(num_channels))
         self.eps = eps
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Perform forward pass for 2D layer normalization.
 
@@ -407,7 +458,8 @@ class MSDeformAttn(nn.Module):
     """
     Multiscale Deformable Attention Module based on Deformable-DETR and PaddleDetection implementations.
 
-    https://github.com/fundamentalvision/Deformable-DETR/blob/main/models/ops/modules/ms_deform_attn.py
+    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.
@@ -419,9 +471,12 @@ class MSDeformAttn(nn.Module):
         attention_weights (nn.Linear): Linear layer for generating attention weights.
         value_proj (nn.Linear): Linear layer for projecting values.
         output_proj (nn.Linear): Linear layer for projecting output.
+
+    References:
+        https://github.com/fundamentalvision/Deformable-DETR/blob/main/models/ops/modules/ms_deform_attn.py
     """
 
-    def __init__(self, d_model=256, n_levels=4, n_heads=8, n_points=4):
+    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.
 
@@ -473,23 +528,31 @@ class MSDeformAttn(nn.Module):
         xavier_uniform_(self.output_proj.weight.data)
         constant_(self.output_proj.bias.data, 0.0)
 
-    def forward(self, query, refer_bbox, value, value_shapes, value_mask=None):
+    def forward(
+        self,
+        query: torch.Tensor,
+        refer_bbox: torch.Tensor,
+        value: torch.Tensor,
+        value_shapes: List,
+        value_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
         """
         Perform forward pass for multiscale deformable attention.
 
-        https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/transformers/deformable_transformer.py
-
         Args:
-            query (torch.Tensor): Tensor with shape [bs, query_length, C].
-            refer_bbox (torch.Tensor): Tensor 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): Tensor with shape [bs, value_length, C].
+            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.
+            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): Tensor with shape [bs, value_length], True for non-padding elements,
-                False for padding elements.
+            value_mask (torch.Tensor, optional): Mask tensor with shape [bs, value_length], True for non-padding
+                elements, False for padding elements.
 
         Returns:
             (torch.Tensor): Output tensor with shape [bs, Length_{query}, C].
+
+        References:
+            https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/transformers/deformable_transformer.py
         """
         bs, len_q = query.shape[:2]
         len_v = value.shape[1]
@@ -521,8 +584,8 @@ class DeformableTransformerDecoderLayer(nn.Module):
     """
     Deformable Transformer Decoder Layer inspired by PaddleDetection and Deformable-DETR implementations.
 
-    https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/transformers/deformable_transformer.py
-    https://github.com/fundamentalvision/Deformable-DETR/blob/main/models/deformable_transformer.py
+    This class implements a single decoder layer with self-attention, cross-attention using multiscale deformable
+    attention, and a feedforward network.
 
     Attributes:
         self_attn (nn.MultiheadAttention): Self-attention module.
@@ -537,9 +600,22 @@ class DeformableTransformerDecoderLayer(nn.Module):
         linear2 (nn.Linear): Second linear layer in the feedforward network.
         dropout4 (nn.Dropout): Dropout after the feedforward network.
         norm3 (nn.LayerNorm): Layer normalization after the feedforward network.
+
+    References:
+        https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/transformers/deformable_transformer.py
+        https://github.com/fundamentalvision/Deformable-DETR/blob/main/models/deformable_transformer.py
     """
 
-    def __init__(self, d_model=256, n_heads=8, d_ffn=1024, dropout=0.0, act=nn.ReLU(), n_levels=4, n_points=4):
+    def __init__(
+        self,
+        d_model: int = 256,
+        n_heads: int = 8,
+        d_ffn: int = 1024,
+        dropout: float = 0.0,
+        act: nn.Module = nn.ReLU(),
+        n_levels: int = 4,
+        n_points: int = 4,
+    ):
         """
         Initialize the DeformableTransformerDecoderLayer with the given parameters.
 
@@ -573,11 +649,11 @@ class DeformableTransformerDecoderLayer(nn.Module):
         self.norm3 = nn.LayerNorm(d_model)
 
     @staticmethod
-    def with_pos_embed(tensor, pos):
+    def with_pos_embed(tensor: torch.Tensor, pos: Optional[torch.Tensor]) -> torch.Tensor:
         """Add positional embeddings to the input tensor, if provided."""
         return tensor if pos is None else tensor + pos
 
-    def forward_ffn(self, tgt):
+    def forward_ffn(self, tgt: torch.Tensor) -> torch.Tensor:
         """
         Perform forward pass through the Feed-Forward Network part of the layer.
 
@@ -591,7 +667,16 @@ class DeformableTransformerDecoderLayer(nn.Module):
         tgt = tgt + self.dropout4(tgt2)
         return self.norm3(tgt)
 
-    def forward(self, embed, refer_bbox, feats, shapes, padding_mask=None, attn_mask=None, query_pos=None):
+    def forward(
+        self,
+        embed: torch.Tensor,
+        refer_bbox: torch.Tensor,
+        feats: torch.Tensor,
+        shapes: List,
+        padding_mask: Optional[torch.Tensor] = None,
+        attn_mask: Optional[torch.Tensor] = None,
+        query_pos: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
         """
         Perform the forward pass through the entire decoder layer.
 
@@ -628,18 +713,22 @@ class DeformableTransformerDecoderLayer(nn.Module):
 
 class DeformableTransformerDecoder(nn.Module):
     """
-    Implementation of Deformable Transformer Decoder based on PaddleDetection.
+    Deformable Transformer Decoder based on PaddleDetection implementation.
 
-    https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/transformers/deformable_transformer.py
+    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.
         num_layers (int): Number of decoder layers.
         hidden_dim (int): Hidden dimension.
         eval_idx (int): Index of the layer to use during evaluation.
+
+    References:
+        https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/transformers/deformable_transformer.py
     """
 
-    def __init__(self, hidden_dim, decoder_layer, num_layers, eval_idx=-1):
+    def __init__(self, hidden_dim: int, decoder_layer: nn.Module, num_layers: int, eval_idx: int = -1):
         """
         Initialize the DeformableTransformerDecoder with the given parameters.
 
@@ -657,15 +746,15 @@ class DeformableTransformerDecoder(nn.Module):
 
     def forward(
         self,
-        embed,  # decoder embeddings
-        refer_bbox,  # anchor
-        feats,  # image features
-        shapes,  # feature shapes
-        bbox_head,
-        score_head,
-        pos_mlp,
-        attn_mask=None,
-        padding_mask=None,
+        embed: torch.Tensor,  # decoder embeddings
+        refer_bbox: torch.Tensor,  # anchor
+        feats: torch.Tensor,  # image features
+        shapes: List,  # feature shapes
+        bbox_head: nn.Module,
+        score_head: nn.Module,
+        pos_mlp: nn.Module,
+        attn_mask: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
     ):
         """
         Perform the forward pass through the entire decoder.