init commit

ultralytics/nn/modules/__init__.py

@@ -0,0 +1,182 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+"""
+Ultralytics neural network modules.
+
+This module provides access to various neural network components used in Ultralytics models, including convolution
+blocks, attention mechanisms, transformer components, and detection/segmentation heads.
+
+Examples:
+    Visualize a module with Netron
+    >>> from ultralytics.nn.modules import Conv
+    >>> import torch
+    >>> import subprocess
+    >>> x = torch.ones(1, 128, 40, 40)
+    >>> m = Conv(128, 128)
+    >>> f = f"{m._get_name()}.onnx"
+    >>> torch.onnx.export(m, x, f)
+    >>> subprocess.run(f"onnxslim {f} {f} && open {f}", shell=True, check=True)  # pip install onnxslim
+"""
+
+from .block import (
+    C1,
+    C2,
+    C2PSA,
+    C3,
+    C3TR,
+    CIB,
+    DFL,
+    ELAN1,
+    PSA,
+    SPP,
+    SPPELAN,
+    SPPF,
+    A2C2f,
+    AConv,
+    ADown,
+    Attention,
+    BNContrastiveHead,
+    Bottleneck,
+    BottleneckCSP,
+    C2f,
+    C2fAttn,
+    C2fCIB,
+    C2fPSA,
+    C3Ghost,
+    C3k2,
+    C3x,
+    CBFuse,
+    CBLinear,
+    ContrastiveHead,
+    GhostBottleneck,
+    HGBlock,
+    HGStem,
+    ImagePoolingAttn,
+    MaxSigmoidAttnBlock,
+    Proto,
+    RepC3,
+    RepNCSPELAN4,
+    RepVGGDW,
+    ResNetLayer,
+    SCDown,
+    TorchVision,
+)
+from .conv import (
+    CBAM,
+    ChannelAttention,
+    Concat,
+    Conv,
+    Conv2,
+    ConvTranspose,
+    DWConv,
+    DWConvTranspose2d,
+    Focus,
+    GhostConv,
+    Index,
+    LightConv,
+    RepConv,
+    SpatialAttention,
+)
+from .head import (
+    OBB,
+    Classify,
+    Detect,
+    LRPCHead,
+    Pose,
+    RTDETRDecoder,
+    Segment,
+    WorldDetect,
+    YOLOEDetect,
+    YOLOESegment,
+    v10Detect,
+)
+from .transformer import (
+    AIFI,
+    MLP,
+    DeformableTransformerDecoder,
+    DeformableTransformerDecoderLayer,
+    LayerNorm2d,
+    MLPBlock,
+    MSDeformAttn,
+    TransformerBlock,
+    TransformerEncoderLayer,
+    TransformerLayer,
+)
+
+__all__ = (
+    "Conv",
+    "Conv2",
+    "LightConv",
+    "RepConv",
+    "DWConv",
+    "DWConvTranspose2d",
+    "ConvTranspose",
+    "Focus",
+    "GhostConv",
+    "ChannelAttention",
+    "SpatialAttention",
+    "CBAM",
+    "Concat",
+    "TransformerLayer",
+    "TransformerBlock",
+    "MLPBlock",
+    "LayerNorm2d",
+    "DFL",
+    "HGBlock",
+    "HGStem",
+    "SPP",
+    "SPPF",
+    "C1",
+    "C2",
+    "C3",
+    "C2f",
+    "C3k2",
+    "SCDown",
+    "C2fPSA",
+    "C2PSA",
+    "C2fAttn",
+    "C3x",
+    "C3TR",
+    "C3Ghost",
+    "GhostBottleneck",
+    "Bottleneck",
+    "BottleneckCSP",
+    "Proto",
+    "Detect",
+    "Segment",
+    "Pose",
+    "Classify",
+    "TransformerEncoderLayer",
+    "RepC3",
+    "RTDETRDecoder",
+    "AIFI",
+    "DeformableTransformerDecoder",
+    "DeformableTransformerDecoderLayer",
+    "MSDeformAttn",
+    "MLP",
+    "ResNetLayer",
+    "OBB",
+    "WorldDetect",
+    "YOLOEDetect",
+    "YOLOESegment",
+    "v10Detect",
+    "LRPCHead",
+    "ImagePoolingAttn",
+    "MaxSigmoidAttnBlock",
+    "ContrastiveHead",
+    "BNContrastiveHead",
+    "RepNCSPELAN4",
+    "ADown",
+    "SPPELAN",
+    "CBFuse",
+    "CBLinear",
+    "AConv",
+    "ELAN1",
+    "RepVGGDW",
+    "CIB",
+    "C2fCIB",
+    "Attention",
+    "PSA",
+    "TorchVision",
+    "Index",
+    "A2C2f",
+)

ultralytics/nn/modules/activation.py

@@ -0,0 +1,56 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+"""Activation modules."""
+
+import torch
+import torch.nn as nn
+
+
+class AGLU(nn.Module):
+    """
+    Unified activation function module from AGLU.
+
+    This class implements a parameterized activation function with learnable parameters lambda and kappa, based on the
+    AGLU (Adaptive Gated Linear Unit) approach.
+
+    Attributes:
+        act (nn.Softplus): Softplus activation function with negative beta.
+        lambd (nn.Parameter): Learnable lambda parameter initialized with uniform distribution.
+        kappa (nn.Parameter): Learnable kappa parameter initialized with uniform distribution.
+
+    Methods:
+        forward: Compute the forward pass of the Unified activation function.
+
+    Examples:
+        >>> import torch
+        >>> m = AGLU()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+        >>> print(output.shape)
+        torch.Size([2])
+
+    References:
+        https://github.com/kostas1515/AGLU
+    """
+
+    def __init__(self, device=None, dtype=None) -> None:
+        """Initialize the Unified activation function with learnable parameters."""
+        super().__init__()
+        self.act = nn.Softplus(beta=-1.0)
+        self.lambd = nn.Parameter(nn.init.uniform_(torch.empty(1, device=device, dtype=dtype)))  # lambda parameter
+        self.kappa = nn.Parameter(nn.init.uniform_(torch.empty(1, device=device, dtype=dtype)))  # kappa parameter
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Apply the Adaptive Gated Linear Unit (AGLU) activation function.
+
+        This forward method implements the AGLU activation function with learnable parameters lambda and kappa.
+        The function applies a transformation that adaptively combines linear and non-linear components.
+
+        Args:
+            x (torch.Tensor): Input tensor to apply the activation function to.
+
+        Returns:
+            (torch.Tensor): Output tensor after applying the AGLU activation function, with the same shape as the input.
+        """
+        lam = torch.clamp(self.lambd, min=0.0001)  # Clamp lambda to avoid division by zero
+        return torch.exp((1 / lam) * self.act((self.kappa * x) - torch.log(lam)))

ultralytics/nn/modules/conv.py

@@ -0,0 +1,714 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+"""Convolution modules."""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+__all__ = (
+    "Conv",
+    "Conv2",
+    "LightConv",
+    "DWConv",
+    "DWConvTranspose2d",
+    "ConvTranspose",
+    "Focus",
+    "GhostConv",
+    "ChannelAttention",
+    "SpatialAttention",
+    "CBAM",
+    "Concat",
+    "RepConv",
+    "Index",
+)
+
+
+def autopad(k, p=None, d=1):  # kernel, padding, dilation
+    """Pad to 'same' shape outputs."""
+    if d > 1:
+        k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k]  # actual kernel-size
+    if p is None:
+        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
+    return p
+
+
+class Conv(nn.Module):
+    """
+    Standard convolution module with batch normalization and activation.
+
+    Attributes:
+        conv (nn.Conv2d): Convolutional layer.
+        bn (nn.BatchNorm2d): Batch normalization layer.
+        act (nn.Module): Activation function layer.
+        default_act (nn.Module): Default activation function (SiLU).
+    """
+
+    default_act = nn.SiLU()  # default activation
+
+    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
+        """
+        Initialize Conv layer with given parameters.
+
+        Args:
+            c1 (int): Number of input channels.
+            c2 (int): Number of output channels.
+            k (int): Kernel size.
+            s (int): Stride.
+            p (int, optional): Padding.
+            g (int): Groups.
+            d (int): Dilation.
+            act (bool | nn.Module): Activation function.
+        """
+        super().__init__()
+        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
+        self.bn = nn.BatchNorm2d(c2)
+        self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
+
+    def forward(self, x):
+        """
+        Apply convolution, batch normalization and activation to input tensor.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor.
+        """
+        return self.act(self.bn(self.conv(x)))
+
+    def forward_fuse(self, x):
+        """
+        Apply convolution and activation without batch normalization.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor.
+        """
+        return self.act(self.conv(x))
+
+
+class Conv2(Conv):
+    """
+    Simplified RepConv module with Conv fusing.
+
+    Attributes:
+        conv (nn.Conv2d): Main 3x3 convolutional layer.
+        cv2 (nn.Conv2d): Additional 1x1 convolutional layer.
+        bn (nn.BatchNorm2d): Batch normalization layer.
+        act (nn.Module): Activation function layer.
+    """
+
+    def __init__(self, c1, c2, k=3, s=1, p=None, g=1, d=1, act=True):
+        """
+        Initialize Conv2 layer with given parameters.
+
+        Args:
+            c1 (int): Number of input channels.
+            c2 (int): Number of output channels.
+            k (int): Kernel size.
+            s (int): Stride.
+            p (int, optional): Padding.
+            g (int): Groups.
+            d (int): Dilation.
+            act (bool | nn.Module): Activation function.
+        """
+        super().__init__(c1, c2, k, s, p, g=g, d=d, act=act)
+        self.cv2 = nn.Conv2d(c1, c2, 1, s, autopad(1, p, d), groups=g, dilation=d, bias=False)  # add 1x1 conv
+
+    def forward(self, x):
+        """
+        Apply convolution, batch normalization and activation to input tensor.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor.
+        """
+        return self.act(self.bn(self.conv(x) + self.cv2(x)))
+
+    def forward_fuse(self, x):
+        """
+        Apply fused convolution, batch normalization and activation to input tensor.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor.
+        """
+        return self.act(self.bn(self.conv(x)))
+
+    def fuse_convs(self):
+        """Fuse parallel convolutions."""
+        w = torch.zeros_like(self.conv.weight.data)
+        i = [x // 2 for x in w.shape[2:]]
+        w[:, :, i[0] : i[0] + 1, i[1] : i[1] + 1] = self.cv2.weight.data.clone()
+        self.conv.weight.data += w
+        self.__delattr__("cv2")
+        self.forward = self.forward_fuse
+
+
+class LightConv(nn.Module):
+    """
+    Light convolution module with 1x1 and depthwise convolutions.
+
+    This implementation is based on the PaddleDetection HGNetV2 backbone.
+
+    Attributes:
+        conv1 (Conv): 1x1 convolution layer.
+        conv2 (DWConv): Depthwise convolution layer.
+    """
+
+    def __init__(self, c1, c2, k=1, act=nn.ReLU()):
+        """
+        Initialize LightConv layer with given parameters.
+
+        Args:
+            c1 (int): Number of input channels.
+            c2 (int): Number of output channels.
+            k (int): Kernel size for depthwise convolution.
+            act (nn.Module): Activation function.
+        """
+        super().__init__()
+        self.conv1 = Conv(c1, c2, 1, act=False)
+        self.conv2 = DWConv(c2, c2, k, act=act)
+
+    def forward(self, x):
+        """
+        Apply 2 convolutions to input tensor.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor.
+        """
+        return self.conv2(self.conv1(x))
+
+
+class DWConv(Conv):
+    """Depth-wise convolution module."""
+
+    def __init__(self, c1, c2, k=1, s=1, d=1, act=True):
+        """
+        Initialize depth-wise convolution with given parameters.
+
+        Args:
+            c1 (int): Number of input channels.
+            c2 (int): Number of output channels.
+            k (int): Kernel size.
+            s (int): Stride.
+            d (int): Dilation.
+            act (bool | nn.Module): Activation function.
+        """
+        super().__init__(c1, c2, k, s, g=math.gcd(c1, c2), d=d, act=act)
+
+
+class DWConvTranspose2d(nn.ConvTranspose2d):
+    """Depth-wise transpose convolution module."""
+
+    def __init__(self, c1, c2, k=1, s=1, p1=0, p2=0):
+        """
+        Initialize depth-wise transpose convolution with given parameters.
+
+        Args:
+            c1 (int): Number of input channels.
+            c2 (int): Number of output channels.
+            k (int): Kernel size.
+            s (int): Stride.
+            p1 (int): Padding.
+            p2 (int): Output padding.
+        """
+        super().__init__(c1, c2, k, s, p1, p2, groups=math.gcd(c1, c2))
+
+
+class ConvTranspose(nn.Module):
+    """
+    Convolution transpose module with optional batch normalization and activation.
+
+    Attributes:
+        conv_transpose (nn.ConvTranspose2d): Transposed convolution layer.
+        bn (nn.BatchNorm2d | nn.Identity): Batch normalization layer.
+        act (nn.Module): Activation function layer.
+        default_act (nn.Module): Default activation function (SiLU).
+    """
+
+    default_act = nn.SiLU()  # default activation
+
+    def __init__(self, c1, c2, k=2, s=2, p=0, bn=True, act=True):
+        """
+        Initialize ConvTranspose layer with given parameters.
+
+        Args:
+            c1 (int): Number of input channels.
+            c2 (int): Number of output channels.
+            k (int): Kernel size.
+            s (int): Stride.
+            p (int): Padding.
+            bn (bool): Use batch normalization.
+            act (bool | nn.Module): Activation function.
+        """
+        super().__init__()
+        self.conv_transpose = nn.ConvTranspose2d(c1, c2, k, s, p, bias=not bn)
+        self.bn = nn.BatchNorm2d(c2) if bn else nn.Identity()
+        self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
+
+    def forward(self, x):
+        """
+        Apply transposed convolution, batch normalization and activation to input.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor.
+        """
+        return self.act(self.bn(self.conv_transpose(x)))
+
+    def forward_fuse(self, x):
+        """
+        Apply activation and convolution transpose operation to input.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor.
+        """
+        return self.act(self.conv_transpose(x))
+
+
+class Focus(nn.Module):
+    """
+    Focus module for concentrating feature information.
+
+    Slices input tensor into 4 parts and concatenates them in the channel dimension.
+
+    Attributes:
+        conv (Conv): Convolution layer.
+    """
+
+    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
+        """
+        Initialize Focus module with given parameters.
+
+        Args:
+            c1 (int): Number of input channels.
+            c2 (int): Number of output channels.
+            k (int): Kernel size.
+            s (int): Stride.
+            p (int, optional): Padding.
+            g (int): Groups.
+            act (bool | nn.Module): Activation function.
+        """
+        super().__init__()
+        self.conv = Conv(c1 * 4, c2, k, s, p, g, act=act)
+        # self.contract = Contract(gain=2)
+
+    def forward(self, x):
+        """
+        Apply Focus operation and convolution to input tensor.
+
+        Input shape is (B, C, W, H) and output shape is (B, 4C, W/2, H/2).
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor.
+        """
+        return self.conv(torch.cat((x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]), 1))
+        # return self.conv(self.contract(x))
+
+
+class GhostConv(nn.Module):
+    """
+    Ghost Convolution module.
+
+    Generates more features with fewer parameters by using cheap operations.
+
+    Attributes:
+        cv1 (Conv): Primary convolution.
+        cv2 (Conv): Cheap operation convolution.
+
+    References:
+        https://github.com/huawei-noah/Efficient-AI-Backbones
+    """
+
+    def __init__(self, c1, c2, k=1, s=1, g=1, act=True):
+        """
+        Initialize Ghost Convolution module with given parameters.
+
+        Args:
+            c1 (int): Number of input channels.
+            c2 (int): Number of output channels.
+            k (int): Kernel size.
+            s (int): Stride.
+            g (int): Groups.
+            act (bool | nn.Module): Activation function.
+        """
+        super().__init__()
+        c_ = c2 // 2  # hidden channels
+        self.cv1 = Conv(c1, c_, k, s, None, g, act=act)
+        self.cv2 = Conv(c_, c_, 5, 1, None, c_, act=act)
+
+    def forward(self, x):
+        """
+        Apply Ghost Convolution to input tensor.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor with concatenated features.
+        """
+        y = self.cv1(x)
+        return torch.cat((y, self.cv2(y)), 1)
+
+
+class RepConv(nn.Module):
+    """
+    RepConv module with training and deploy modes.
+
+    This module is used in RT-DETR and can fuse convolutions during inference for efficiency.
+
+    Attributes:
+        conv1 (Conv): 3x3 convolution.
+        conv2 (Conv): 1x1 convolution.
+        bn (nn.BatchNorm2d, optional): Batch normalization for identity branch.
+        act (nn.Module): Activation function.
+        default_act (nn.Module): Default activation function (SiLU).
+
+    References:
+        https://github.com/DingXiaoH/RepVGG/blob/main/repvgg.py
+    """
+
+    default_act = nn.SiLU()  # default activation
+
+    def __init__(self, c1, c2, k=3, s=1, p=1, g=1, d=1, act=True, bn=False, deploy=False):
+        """
+        Initialize RepConv module with given parameters.
+
+        Args:
+            c1 (int): Number of input channels.
+            c2 (int): Number of output channels.
+            k (int): Kernel size.
+            s (int): Stride.
+            p (int): Padding.
+            g (int): Groups.
+            d (int): Dilation.
+            act (bool | nn.Module): Activation function.
+            bn (bool): Use batch normalization for identity branch.
+            deploy (bool): Deploy mode for inference.
+        """
+        super().__init__()
+        assert k == 3 and p == 1
+        self.g = g
+        self.c1 = c1
+        self.c2 = c2
+        self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
+
+        self.bn = nn.BatchNorm2d(num_features=c1) if bn and c2 == c1 and s == 1 else None
+        self.conv1 = Conv(c1, c2, k, s, p=p, g=g, act=False)
+        self.conv2 = Conv(c1, c2, 1, s, p=(p - k // 2), g=g, act=False)
+
+    def forward_fuse(self, x):
+        """
+        Forward pass for deploy mode.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor.
+        """
+        return self.act(self.conv(x))
+
+    def forward(self, x):
+        """
+        Forward pass for training mode.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor.
+        """
+        id_out = 0 if self.bn is None else self.bn(x)
+        return self.act(self.conv1(x) + self.conv2(x) + id_out)
+
+    def get_equivalent_kernel_bias(self):
+        """
+        Calculate equivalent kernel and bias by fusing convolutions.
+
+        Returns:
+            (torch.Tensor): Equivalent kernel
+            (torch.Tensor): Equivalent bias
+        """
+        kernel3x3, bias3x3 = self._fuse_bn_tensor(self.conv1)
+        kernel1x1, bias1x1 = self._fuse_bn_tensor(self.conv2)
+        kernelid, biasid = self._fuse_bn_tensor(self.bn)
+        return kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1) + kernelid, bias3x3 + bias1x1 + biasid
+
+    @staticmethod
+    def _pad_1x1_to_3x3_tensor(kernel1x1):
+        """
+        Pad a 1x1 kernel to 3x3 size.
+
+        Args:
+            kernel1x1 (torch.Tensor): 1x1 convolution kernel.
+
+        Returns:
+            (torch.Tensor): Padded 3x3 kernel.
+        """
+        if kernel1x1 is None:
+            return 0
+        else:
+            return torch.nn.functional.pad(kernel1x1, [1, 1, 1, 1])
+
+    def _fuse_bn_tensor(self, branch):
+        """
+        Fuse batch normalization with convolution weights.
+
+        Args:
+            branch (Conv | nn.BatchNorm2d | None): Branch to fuse.
+
+        Returns:
+            kernel (torch.Tensor): Fused kernel.
+            bias (torch.Tensor): Fused bias.
+        """
+        if branch is None:
+            return 0, 0
+        if isinstance(branch, Conv):
+            kernel = branch.conv.weight
+            running_mean = branch.bn.running_mean
+            running_var = branch.bn.running_var
+            gamma = branch.bn.weight
+            beta = branch.bn.bias
+            eps = branch.bn.eps
+        elif isinstance(branch, nn.BatchNorm2d):
+            if not hasattr(self, "id_tensor"):
+                input_dim = self.c1 // self.g
+                kernel_value = np.zeros((self.c1, input_dim, 3, 3), dtype=np.float32)
+                for i in range(self.c1):
+                    kernel_value[i, i % input_dim, 1, 1] = 1
+                self.id_tensor = torch.from_numpy(kernel_value).to(branch.weight.device)
+            kernel = self.id_tensor
+            running_mean = branch.running_mean
+            running_var = branch.running_var
+            gamma = branch.weight
+            beta = branch.bias
+            eps = branch.eps
+        std = (running_var + eps).sqrt()
+        t = (gamma / std).reshape(-1, 1, 1, 1)
+        return kernel * t, beta - running_mean * gamma / std
+
+    def fuse_convs(self):
+        """Fuse convolutions for inference by creating a single equivalent convolution."""
+        if hasattr(self, "conv"):
+            return
+        kernel, bias = self.get_equivalent_kernel_bias()
+        self.conv = nn.Conv2d(
+            in_channels=self.conv1.conv.in_channels,
+            out_channels=self.conv1.conv.out_channels,
+            kernel_size=self.conv1.conv.kernel_size,
+            stride=self.conv1.conv.stride,
+            padding=self.conv1.conv.padding,
+            dilation=self.conv1.conv.dilation,
+            groups=self.conv1.conv.groups,
+            bias=True,
+        ).requires_grad_(False)
+        self.conv.weight.data = kernel
+        self.conv.bias.data = bias
+        for para in self.parameters():
+            para.detach_()
+        self.__delattr__("conv1")
+        self.__delattr__("conv2")
+        if hasattr(self, "nm"):
+            self.__delattr__("nm")
+        if hasattr(self, "bn"):
+            self.__delattr__("bn")
+        if hasattr(self, "id_tensor"):
+            self.__delattr__("id_tensor")
+
+
+class ChannelAttention(nn.Module):
+    """
+    Channel-attention module for feature recalibration.
+
+    Applies attention weights to channels based on global average pooling.
+
+    Attributes:
+        pool (nn.AdaptiveAvgPool2d): Global average pooling.
+        fc (nn.Conv2d): Fully connected layer implemented as 1x1 convolution.
+        act (nn.Sigmoid): Sigmoid activation for attention weights.
+
+    References:
+        https://github.com/open-mmlab/mmdetection/tree/v3.0.0rc1/configs/rtmdet
+    """
+
+    def __init__(self, channels: int) -> None:
+        """
+        Initialize Channel-attention module.
+
+        Args:
+            channels (int): Number of input channels.
+        """
+        super().__init__()
+        self.pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Conv2d(channels, channels, 1, 1, 0, bias=True)
+        self.act = nn.Sigmoid()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Apply channel attention to input tensor.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Channel-attended output tensor.
+        """
+        return x * self.act(self.fc(self.pool(x)))
+
+
+class SpatialAttention(nn.Module):
+    """
+    Spatial-attention module for feature recalibration.
+
+    Applies attention weights to spatial dimensions based on channel statistics.
+
+    Attributes:
+        cv1 (nn.Conv2d): Convolution layer for spatial attention.
+        act (nn.Sigmoid): Sigmoid activation for attention weights.
+    """
+
+    def __init__(self, kernel_size=7):
+        """
+        Initialize Spatial-attention module.
+
+        Args:
+            kernel_size (int): Size of the convolutional kernel (3 or 7).
+        """
+        super().__init__()
+        assert kernel_size in {3, 7}, "kernel size must be 3 or 7"
+        padding = 3 if kernel_size == 7 else 1
+        self.cv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
+        self.act = nn.Sigmoid()
+
+    def forward(self, x):
+        """
+        Apply spatial attention to input tensor.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Spatial-attended output tensor.
+        """
+        return x * self.act(self.cv1(torch.cat([torch.mean(x, 1, keepdim=True), torch.max(x, 1, keepdim=True)[0]], 1)))
+
+
+class CBAM(nn.Module):
+    """
+    Convolutional Block Attention Module.
+
+    Combines channel and spatial attention mechanisms for comprehensive feature refinement.
+
+    Attributes:
+        channel_attention (ChannelAttention): Channel attention module.
+        spatial_attention (SpatialAttention): Spatial attention module.
+    """
+
+    def __init__(self, c1, kernel_size=7):
+        """
+        Initialize CBAM with given parameters.
+
+        Args:
+            c1 (int): Number of input channels.
+            kernel_size (int): Size of the convolutional kernel for spatial attention.
+        """
+        super().__init__()
+        self.channel_attention = ChannelAttention(c1)
+        self.spatial_attention = SpatialAttention(kernel_size)
+
+    def forward(self, x):
+        """
+        Apply channel and spatial attention sequentially to input tensor.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Attended output tensor.
+        """
+        return self.spatial_attention(self.channel_attention(x))
+
+
+class Concat(nn.Module):
+    """
+    Concatenate a list of tensors along specified dimension.
+
+    Attributes:
+        d (int): Dimension along which to concatenate tensors.
+    """
+
+    def __init__(self, dimension=1):
+        """
+        Initialize Concat module.
+
+        Args:
+            dimension (int): Dimension along which to concatenate tensors.
+        """
+        super().__init__()
+        self.d = dimension
+
+    def forward(self, x: list[torch.Tensor]):
+        """
+        Concatenate input tensors along specified dimension.
+
+        Args:
+            x (list[torch.Tensor]): List of input tensors.
+
+        Returns:
+            (torch.Tensor): Concatenated tensor.
+        """
+        return torch.cat(x, self.d)
+
+
+class Index(nn.Module):
+    """
+    Returns a particular index of the input.
+
+    Attributes:
+        index (int): Index to select from input.
+    """
+
+    def __init__(self, index=0):
+        """
+        Initialize Index module.
+
+        Args:
+            index (int): Index to select from input.
+        """
+        super().__init__()
+        self.index = index
+
+    def forward(self, x: list[torch.Tensor]):
+        """
+        Select and return a particular index from input.
+
+        Args:
+            x (list[torch.Tensor]): List of input tensors.
+
+        Returns:
+            (torch.Tensor): Selected tensor.
+        """
+        return x[self.index]

ultralytics/nn/modules/transformer.py

@@ -0,0 +1,805 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+"""Transformer modules."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+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",
+    "DeformableTransformerDecoder",
+    "DeformableTransformerDecoderLayer",
+    "MSDeformAttn",
+    "MLP",
+)
+
+
+class TransformerEncoderLayer(nn.Module):
+    """
+    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.
+        fc1 (nn.Linear): First linear layer in the feedforward network.
+        fc2 (nn.Linear): Second linear layer in the feedforward network.
+        norm1 (nn.LayerNorm): Layer normalization after attention.
+        norm2 (nn.LayerNorm): Layer normalization after feedforward network.
+        dropout (nn.Dropout): Dropout layer for the feedforward network.
+        dropout1 (nn.Dropout): Dropout layer after attention.
+        dropout2 (nn.Dropout): Dropout layer after feedforward network.
+        act (nn.Module): Activation function.
+        normalize_before (bool): Whether to apply normalization before attention and feedforward.
+    """
+
+    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.
+
+        Args:
+            c1 (int): Input dimension.
+            cm (int): Hidden dimension in the feedforward network.
+            num_heads (int): Number of attention heads.
+            dropout (float): Dropout probability.
+            act (nn.Module): Activation function.
+            normalize_before (bool): Whether to apply normalization before attention and feedforward.
+        """
+        super().__init__()
+        from ...utils.torch_utils import TORCH_1_9
+
+        if not TORCH_1_9:
+            raise ModuleNotFoundError(
+                "TransformerEncoderLayer() requires torch>=1.9 to use nn.MultiheadAttention(batch_first=True)."
+            )
+        self.ma = nn.MultiheadAttention(c1, num_heads, dropout=dropout, batch_first=True)
+        # Implementation of Feedforward model
+        self.fc1 = nn.Linear(c1, cm)
+        self.fc2 = nn.Linear(cm, c1)
+
+        self.norm1 = nn.LayerNorm(c1)
+        self.norm2 = nn.LayerNorm(c1)
+        self.dropout = nn.Dropout(dropout)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+        self.act = act
+        self.normalize_before = normalize_before
+
+    @staticmethod
+    def with_pos_embed(tensor: torch.Tensor, pos: torch.Tensor | None = 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: torch.Tensor,
+        src_mask: torch.Tensor | None = None,
+        src_key_padding_mask: torch.Tensor | None = None,
+        pos: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        Perform forward pass with post-normalization.
+
+        Args:
+            src (torch.Tensor): Input tensor.
+            src_mask (torch.Tensor, optional): Mask for the src sequence.
+            src_key_padding_mask (torch.Tensor, optional): Mask for the src keys per batch.
+            pos (torch.Tensor, optional): Positional encoding.
+
+        Returns:
+            (torch.Tensor): Output tensor after attention and feedforward.
+        """
+        q = k = self.with_pos_embed(src, pos)
+        src2 = self.ma(q, k, value=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0]
+        src = src + self.dropout1(src2)
+        src = self.norm1(src)
+        src2 = self.fc2(self.dropout(self.act(self.fc1(src))))
+        src = src + self.dropout2(src2)
+        return self.norm2(src)
+
+    def forward_pre(
+        self,
+        src: torch.Tensor,
+        src_mask: torch.Tensor | None = None,
+        src_key_padding_mask: torch.Tensor | None = None,
+        pos: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        Perform forward pass with pre-normalization.
+
+        Args:
+            src (torch.Tensor): Input tensor.
+            src_mask (torch.Tensor, optional): Mask for the src sequence.
+            src_key_padding_mask (torch.Tensor, optional): Mask for the src keys per batch.
+            pos (torch.Tensor, optional): Positional encoding.
+
+        Returns:
+            (torch.Tensor): Output tensor after attention and feedforward.
+        """
+        src2 = self.norm1(src)
+        q = k = self.with_pos_embed(src2, pos)
+        src2 = self.ma(q, k, value=src2, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0]
+        src = src + self.dropout1(src2)
+        src2 = self.norm2(src)
+        src2 = self.fc2(self.dropout(self.act(self.fc1(src2))))
+        return src + self.dropout2(src2)
+
+    def forward(
+        self,
+        src: torch.Tensor,
+        src_mask: torch.Tensor | None = None,
+        src_key_padding_mask: torch.Tensor | None = None,
+        pos: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        Forward propagate the input through the encoder module.
+
+        Args:
+            src (torch.Tensor): Input tensor.
+            src_mask (torch.Tensor, optional): Mask for the src sequence.
+            src_key_padding_mask (torch.Tensor, optional): Mask for the src keys per batch.
+            pos (torch.Tensor, optional): Positional encoding.
+
+        Returns:
+            (torch.Tensor): Output tensor after transformer encoder layer.
+        """
+        if self.normalize_before:
+            return self.forward_pre(src, src_mask, src_key_padding_mask, pos)
+        return self.forward_post(src, src_mask, src_key_padding_mask, pos)
+
+
+class AIFI(TransformerEncoderLayer):
+    """
+    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.
+    """
+
+    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.
+
+        Args:
+            c1 (int): Input dimension.
+            cm (int): Hidden dimension in the feedforward network.
+            num_heads (int): Number of attention heads.
+            dropout (float): Dropout probability.
+            act (nn.Module): Activation function.
+            normalize_before (bool): Whether to apply normalization before attention and feedforward.
+        """
+        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.
+
+        Args:
+            x (torch.Tensor): Input tensor with shape [B, C, H, W].
+
+        Returns:
+            (torch.Tensor): Output tensor with shape [B, C, H, W].
+        """
+        c, h, w = x.shape[1:]
+        pos_embed = self.build_2d_sincos_position_embedding(w, h, c)
+        # Flatten [B, C, H, W] to [B, HxW, C]
+        x = super().forward(x.flatten(2).permute(0, 2, 1), pos=pos_embed.to(device=x.device, dtype=x.dtype))
+        return x.permute(0, 2, 1).view([-1, c, h, w]).contiguous()
+
+    @staticmethod
+    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.
+
+        Args:
+            w (int): Width of the feature map.
+            h (int): Height of the feature map.
+            embed_dim (int): Embedding dimension.
+            temperature (float): Temperature for the sine/cosine functions.
+
+        Returns:
+            (torch.Tensor): Position embedding with shape [1, embed_dim, h*w].
+        """
+        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") 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)
+
+        out_w = grid_w.flatten()[..., None] @ omega[None]
+        out_h = grid_h.flatten()[..., None] @ omega[None]
+
+        return torch.cat([torch.sin(out_w), torch.cos(out_w), torch.sin(out_h), torch.cos(out_h)], 1)[None]
+
+
+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.
+
+        Args:
+            c (int): Input and output channel dimension.
+            num_heads (int): Number of attention heads.
+        """
+        super().__init__()
+        self.q = nn.Linear(c, c, bias=False)
+        self.k = nn.Linear(c, c, bias=False)
+        self.v = nn.Linear(c, c, bias=False)
+        self.ma = nn.MultiheadAttention(embed_dim=c, num_heads=num_heads)
+        self.fc1 = nn.Linear(c, c, bias=False)
+        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.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor after transformer layer.
+        """
+        x = self.ma(self.q(x), self.k(x), self.v(x))[0] + x
+        return self.fc2(self.fc1(x)) + x
+
+
+class TransformerBlock(nn.Module):
+    """
+    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.
+        linear (nn.Linear): Learnable position embedding.
+        tr (nn.Sequential): Sequential container of transformer layers.
+        c2 (int): Output channel dimension.
+    """
+
+    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.
+
+        Args:
+            c1 (int): Input channel dimension.
+            c2 (int): Output channel dimension.
+            num_heads (int): Number of attention heads.
+            num_layers (int): Number of transformer layers.
+        """
+        super().__init__()
+        self.conv = None
+        if c1 != c2:
+            self.conv = Conv(c1, c2)
+        self.linear = nn.Linear(c2, c2)  # learnable position embedding
+        self.tr = nn.Sequential(*(TransformerLayer(c2, num_heads) for _ in range(num_layers)))
+        self.c2 = c2
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward propagate the input through the transformer block.
+
+        Args:
+            x (torch.Tensor): Input tensor with shape [b, c1, w, h].
+
+        Returns:
+            (torch.Tensor): Output tensor with shape [b, c2, w, h].
+        """
+        if self.conv is not None:
+            x = self.conv(x)
+        b, _, w, h = 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)
+
+
+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.
+
+        Args:
+            embedding_dim (int): Input and output dimension.
+            mlp_dim (int): Hidden dimension.
+            act (nn.Module): Activation function.
+        """
+        super().__init__()
+        self.lin1 = nn.Linear(embedding_dim, mlp_dim)
+        self.lin2 = nn.Linear(mlp_dim, embedding_dim)
+        self.act = act()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass for the MLPBlock.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor after MLP block.
+        """
+        return self.lin2(self.act(self.lin1(x)))
+
+
+class MLP(nn.Module):
+    """
+    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.
+        layers (nn.ModuleList): List of linear layers.
+        sigmoid (bool): Whether to apply sigmoid to the output.
+        act (nn.Module): Activation function.
+    """
+
+    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.
+
+        Args:
+            input_dim (int): Input dimension.
+            hidden_dim (int): Hidden dimension.
+            output_dim (int): Output dimension.
+            num_layers (int): Number of layers.
+            act (nn.Module): Activation function.
+            sigmoid (bool): Whether to apply sigmoid to 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.sigmoid = sigmoid
+        self.act = act()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass for the entire MLP.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor after MLP.
+        """
+        for i, layer in enumerate(self.layers):
+            x = getattr(self, "act", nn.ReLU())(layer(x)) if i < self.num_layers - 1 else layer(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.
+
+    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: int, eps: float = 1e-6):
+        """
+        Initialize LayerNorm2d with the given parameters.
+
+        Args:
+            num_channels (int): Number of channels in the input.
+            eps (float): Small constant for numerical stability.
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(num_channels))
+        self.bias = nn.Parameter(torch.zeros(num_channels))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Perform forward pass for 2D layer normalization.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Normalized output tensor.
+        """
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.eps)
+        return self.weight[:, None, None] * x + self.bias[:, None, None]
+
+
+class MSDeformAttn(nn.Module):
+    """
+    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.
+
+    Attributes:
+        im2col_step (int): Step size for im2col operations.
+        d_model (int): Model dimension.
+        n_levels (int): Number of feature levels.
+        n_heads (int): Number of attention heads.
+        n_points (int): Number of sampling points per attention head per feature level.
+        sampling_offsets (nn.Linear): Linear layer for generating sampling offsets.
+        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: int = 256, n_levels: int = 4, n_heads: int = 8, n_points: int = 4):
+        """
+        Initialize MSDeformAttn with the given parameters.
+
+        Args:
+            d_model (int): Model dimension.
+            n_levels (int): Number of feature levels.
+            n_heads (int): Number of attention heads.
+            n_points (int): Number of sampling points per attention head per feature level.
+        """
+        super().__init__()
+        if d_model % n_heads != 0:
+            raise ValueError(f"d_model must be divisible by n_heads, but got {d_model} and {n_heads}")
+        _d_per_head = d_model // n_heads
+        # Better to set _d_per_head to a power of 2 which is more efficient in a CUDA implementation
+        assert _d_per_head * n_heads == d_model, "`d_model` must be divisible by `n_heads`"
+
+        self.im2col_step = 64
+
+        self.d_model = d_model
+        self.n_levels = n_levels
+        self.n_heads = n_heads
+        self.n_points = n_points
+
+        self.sampling_offsets = nn.Linear(d_model, n_heads * n_levels * n_points * 2)
+        self.attention_weights = nn.Linear(d_model, n_heads * n_levels * n_points)
+        self.value_proj = nn.Linear(d_model, d_model)
+        self.output_proj = nn.Linear(d_model, d_model)
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        """Reset module parameters."""
+        constant_(self.sampling_offsets.weight.data, 0.0)
+        thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads)
+        grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
+        grid_init = (
+            (grid_init / grid_init.abs().max(-1, keepdim=True)[0])
+            .view(self.n_heads, 1, 1, 2)
+            .repeat(1, self.n_levels, self.n_points, 1)
+        )
+        for i in range(self.n_points):
+            grid_init[:, :, i, :] *= i + 1
+        with torch.no_grad():
+            self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))
+        constant_(self.attention_weights.weight.data, 0.0)
+        constant_(self.attention_weights.bias.data, 0.0)
+        xavier_uniform_(self.value_proj.weight.data)
+        constant_(self.value_proj.bias.data, 0.0)
+        xavier_uniform_(self.output_proj.weight.data)
+        constant_(self.output_proj.bias.data, 0.0)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        refer_bbox: torch.Tensor,
+        value: torch.Tensor,
+        value_shapes: list,
+        value_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        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.
+            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
+                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]
+        assert sum(s[0] * s[1] for s in value_shapes) == len_v
+
+        value = self.value_proj(value)
+        if value_mask is not None:
+            value = value.masked_fill(value_mask[..., None], float(0))
+        value = value.view(bs, len_v, self.n_heads, self.d_model // self.n_heads)
+        sampling_offsets = self.sampling_offsets(query).view(bs, len_q, self.n_heads, self.n_levels, self.n_points, 2)
+        attention_weights = self.attention_weights(query).view(bs, len_q, self.n_heads, self.n_levels * self.n_points)
+        attention_weights = F.softmax(attention_weights, -1).view(bs, len_q, self.n_heads, self.n_levels, self.n_points)
+        # N, Len_q, n_heads, n_levels, n_points, 2
+        num_points = refer_bbox.shape[-1]
+        if num_points == 2:
+            offset_normalizer = torch.as_tensor(value_shapes, dtype=query.dtype, device=query.device).flip(-1)
+            add = sampling_offsets / offset_normalizer[None, None, None, :, None, :]
+            sampling_locations = refer_bbox[:, :, None, :, None, :] + add
+        elif num_points == 4:
+            add = sampling_offsets / self.n_points * refer_bbox[:, :, None, :, None, 2:] * 0.5
+            sampling_locations = refer_bbox[:, :, None, :, None, :2] + add
+        else:
+            raise ValueError(f"Last dim of reference_points must be 2 or 4, but got {num_points}.")
+        output = multi_scale_deformable_attn_pytorch(value, value_shapes, sampling_locations, attention_weights)
+        return self.output_proj(output)
+
+
+class DeformableTransformerDecoderLayer(nn.Module):
+    """
+    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.
+
+    Attributes:
+        self_attn (nn.MultiheadAttention): Self-attention module.
+        dropout1 (nn.Dropout): Dropout after self-attention.
+        norm1 (nn.LayerNorm): Layer normalization after self-attention.
+        cross_attn (MSDeformAttn): Cross-attention module.
+        dropout2 (nn.Dropout): Dropout after cross-attention.
+        norm2 (nn.LayerNorm): Layer normalization after cross-attention.
+        linear1 (nn.Linear): First linear layer in the feedforward network.
+        act (nn.Module): Activation function.
+        dropout3 (nn.Dropout): Dropout in the feedforward network.
+        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: 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.
+
+        Args:
+            d_model (int): Model dimension.
+            n_heads (int): Number of attention heads.
+            d_ffn (int): Dimension of the feedforward network.
+            dropout (float): Dropout probability.
+            act (nn.Module): Activation function.
+            n_levels (int): Number of feature levels.
+            n_points (int): Number of sampling points.
+        """
+        super().__init__()
+
+        # Self attention
+        self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout)
+        self.dropout1 = nn.Dropout(dropout)
+        self.norm1 = nn.LayerNorm(d_model)
+
+        # Cross attention
+        self.cross_attn = MSDeformAttn(d_model, n_levels, n_heads, n_points)
+        self.dropout2 = nn.Dropout(dropout)
+        self.norm2 = nn.LayerNorm(d_model)
+
+        # FFN
+        self.linear1 = nn.Linear(d_model, d_ffn)
+        self.act = act
+        self.dropout3 = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(d_ffn, d_model)
+        self.dropout4 = nn.Dropout(dropout)
+        self.norm3 = nn.LayerNorm(d_model)
+
+    @staticmethod
+    def with_pos_embed(tensor: torch.Tensor, pos: torch.Tensor | None) -> 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: torch.Tensor) -> torch.Tensor:
+        """
+        Perform forward pass through the Feed-Forward Network part of the layer.
+
+        Args:
+            tgt (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor after FFN.
+        """
+        tgt2 = self.linear2(self.dropout3(self.act(self.linear1(tgt))))
+        tgt = tgt + self.dropout4(tgt2)
+        return self.norm3(tgt)
+
+    def forward(
+        self,
+        embed: torch.Tensor,
+        refer_bbox: torch.Tensor,
+        feats: torch.Tensor,
+        shapes: list,
+        padding_mask: torch.Tensor | None = None,
+        attn_mask: torch.Tensor | None = None,
+        query_pos: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        Perform the forward pass through the entire decoder layer.
+
+        Args:
+            embed (torch.Tensor): Input embeddings.
+            refer_bbox (torch.Tensor): Reference bounding boxes.
+            feats (torch.Tensor): Feature maps.
+            shapes (list): Feature shapes.
+            padding_mask (torch.Tensor, optional): Padding mask.
+            attn_mask (torch.Tensor, optional): Attention mask.
+            query_pos (torch.Tensor, optional): Query position embeddings.
+
+        Returns:
+            (torch.Tensor): Output tensor after decoder layer.
+        """
+        # Self attention
+        q = k = self.with_pos_embed(embed, query_pos)
+        tgt = self.self_attn(q.transpose(0, 1), k.transpose(0, 1), embed.transpose(0, 1), attn_mask=attn_mask)[
+            0
+        ].transpose(0, 1)
+        embed = embed + self.dropout1(tgt)
+        embed = self.norm1(embed)
+
+        # Cross attention
+        tgt = self.cross_attn(
+            self.with_pos_embed(embed, query_pos), refer_bbox.unsqueeze(2), feats, shapes, padding_mask
+        )
+        embed = embed + self.dropout2(tgt)
+        embed = self.norm2(embed)
+
+        # FFN
+        return self.forward_ffn(embed)
+
+
+class DeformableTransformerDecoder(nn.Module):
+    """
+    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.
+
+    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: int, decoder_layer: nn.Module, num_layers: int, eval_idx: int = -1):
+        """
+        Initialize the DeformableTransformerDecoder with the given parameters.
+
+        Args:
+            hidden_dim (int): Hidden dimension.
+            decoder_layer (nn.Module): Decoder layer module.
+            num_layers (int): Number of decoder layers.
+            eval_idx (int): Index of the layer to use during evaluation.
+        """
+        super().__init__()
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.hidden_dim = hidden_dim
+        self.eval_idx = eval_idx if eval_idx >= 0 else num_layers + eval_idx
+
+    def forward(
+        self,
+        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: torch.Tensor | None = None,
+        padding_mask: torch.Tensor | None = None,
+    ):
+        """
+        Perform the forward pass through the entire decoder.
+
+        Args:
+            embed (torch.Tensor): Decoder embeddings.
+            refer_bbox (torch.Tensor): Reference bounding boxes.
+            feats (torch.Tensor): Image features.
+            shapes (list): Feature shapes.
+            bbox_head (nn.Module): Bounding box prediction head.
+            score_head (nn.Module): Score prediction head.
+            pos_mlp (nn.Module): Position MLP.
+            attn_mask (torch.Tensor, optional): Attention mask.
+            padding_mask (torch.Tensor, optional): Padding mask.
+
+        Returns:
+            dec_bboxes (torch.Tensor): Decoded bounding boxes.
+            dec_cls (torch.Tensor): Decoded classification scores.
+        """
+        output = embed
+        dec_bboxes = []
+        dec_cls = []
+        last_refined_bbox = None
+        refer_bbox = refer_bbox.sigmoid()
+        for i, layer in enumerate(self.layers):
+            output = layer(output, refer_bbox, feats, shapes, padding_mask, attn_mask, pos_mlp(refer_bbox))
+
+            bbox = bbox_head[i](output)
+            refined_bbox = torch.sigmoid(bbox + inverse_sigmoid(refer_bbox))
+
+            if self.training:
+                dec_cls.append(score_head[i](output))
+                if i == 0:
+                    dec_bboxes.append(refined_bbox)
+                else:
+                    dec_bboxes.append(torch.sigmoid(bbox + inverse_sigmoid(last_refined_bbox)))
+            elif i == self.eval_idx:
+                dec_cls.append(score_head[i](output))
+                dec_bboxes.append(refined_bbox)
+                break
+
+            last_refined_bbox = refined_bbox
+            refer_bbox = refined_bbox.detach() if self.training else refined_bbox
+
+        return torch.stack(dec_bboxes), torch.stack(dec_cls)

ultralytics/nn/modules/utils.py

@@ -0,0 +1,164 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+import copy
+import math
+
+import numpy as np
+import torch
+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"
+
+
+def _get_clones(module, n):
+    """
+    Create a list of cloned modules from the given module.
+
+    Args:
+        module (nn.Module): The module to be cloned.
+        n (int): Number of clones to create.
+
+    Returns:
+        (nn.ModuleList): A ModuleList containing n clones of the input module.
+
+    Examples:
+        >>> import torch.nn as nn
+        >>> layer = nn.Linear(10, 10)
+        >>> clones = _get_clones(layer, 3)
+        >>> len(clones)
+        3
+    """
+    return nn.ModuleList([copy.deepcopy(module) for _ in range(n)])
+
+
+def bias_init_with_prob(prior_prob=0.01):
+    """
+    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.
+
+    Args:
+        prior_prob (float, optional): Prior probability for bias initialization.
+
+    Returns:
+        (float): Bias initialization value calculated from the prior probability.
+
+    Examples:
+        >>> bias = bias_init_with_prob(0.01)
+        >>> print(f"Bias initialization value: {bias:.4f}")
+        Bias initialization value: -4.5951
+    """
+    return float(-np.log((1 - prior_prob) / prior_prob))  # return bias_init
+
+
+def linear_init(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.
+
+    Args:
+        module (nn.Module): Linear module to initialize.
+
+    Returns:
+        (nn.Module): The initialized module.
+
+    Examples:
+        >>> import torch.nn as nn
+        >>> linear = nn.Linear(10, 5)
+        >>> initialized_linear = linear_init(linear)
+    """
+    bound = 1 / math.sqrt(module.weight.shape[0])
+    uniform_(module.weight, -bound, bound)
+    if hasattr(module, "bias") and module.bias is not None:
+        uniform_(module.bias, -bound, bound)
+
+
+def inverse_sigmoid(x, eps=1e-5):
+    """
+    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.
+
+    Args:
+        x (torch.Tensor): Input tensor with values in range [0, 1].
+        eps (float, optional): Small epsilon value to prevent numerical instability.
+
+    Returns:
+        (torch.Tensor): Tensor after applying the inverse sigmoid function.
+
+    Examples:
+        >>> x = torch.tensor([0.2, 0.5, 0.8])
+        >>> inverse_sigmoid(x)
+        tensor([-1.3863,  0.0000,  1.3863])
+    """
+    x = x.clamp(min=0, max=1)
+    x1 = x.clamp(min=eps)
+    x2 = (1 - x).clamp(min=eps)
+    return torch.log(x1 / x2)
+
+
+def multi_scale_deformable_attn_pytorch(
+    value: torch.Tensor,
+    value_spatial_shapes: torch.Tensor,
+    sampling_locations: torch.Tensor,
+    attention_weights: torch.Tensor,
+) -> torch.Tensor:
+    """
+    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.
+
+    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).
+
+    Returns:
+        (torch.Tensor): The output tensor with shape (bs, num_queries, embed_dims).
+
+    References:
+        https://github.com/IDEA-Research/detrex/blob/main/detrex/layers/multi_scale_deform_attn.py
+    """
+    bs, _, num_heads, embed_dims = value.shape
+    _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape
+    value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1)
+    sampling_grids = 2 * sampling_locations - 1
+    sampling_value_list = []
+    for level, (H_, W_) in enumerate(value_spatial_shapes):
+        # bs, H_*W_, num_heads, embed_dims ->
+        # bs, H_*W_, num_heads*embed_dims ->
+        # bs, num_heads*embed_dims, H_*W_ ->
+        # bs*num_heads, embed_dims, H_, W_
+        value_l_ = value_list[level].flatten(2).transpose(1, 2).reshape(bs * num_heads, embed_dims, H_, W_)
+        # bs, num_queries, num_heads, num_points, 2 ->
+        # bs, num_heads, num_queries, num_points, 2 ->
+        # bs*num_heads, num_queries, num_points, 2
+        sampling_grid_l_ = sampling_grids[:, :, :, level].transpose(1, 2).flatten(0, 1)
+        # bs*num_heads, embed_dims, num_queries, num_points
+        sampling_value_l_ = F.grid_sample(
+            value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False
+        )
+        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)
+    attention_weights = attention_weights.transpose(1, 2).reshape(
+        bs * num_heads, 1, num_queries, num_levels * num_points
+    )
+    output = (
+        (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights)
+        .sum(-1)
+        .view(bs, num_heads * embed_dims, num_queries)
+    )
+    return output.transpose(1, 2).contiguous()