dgenerate-ultralytics-headless 8.3.137__py3-none-any.whl → 8.3.224__py3-none-any.whl

ultralytics/nn/modules/conv.py

@@ -1,6 +1,8 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 """Convolution modules."""
 
+from __future__ import annotations
+
 import math
 
 import numpy as np
@@ -8,20 +10,20 @@ import torch
 import torch.nn as nn
 
 __all__ = (
+    "CBAM",
+    "ChannelAttention",
+    "Concat",
     "Conv",
     "Conv2",
-    "LightConv",
+    "ConvTranspose",
     "DWConv",
     "DWConvTranspose2d",
-    "ConvTranspose",
     "Focus",
     "GhostConv",
-    "ChannelAttention",
-    "SpatialAttention",
-    "CBAM",
-    "Concat",
-    "RepConv",
     "Index",
+    "LightConv",
+    "RepConv",
+    "SpatialAttention",
 )
 
 
@@ -35,8 +37,7 @@ def autopad(k, p=None, d=1):  # kernel, padding, dilation
 
 
 class Conv(nn.Module):
-    """
-    Standard convolution module with batch normalization and activation.
+    """Standard convolution module with batch normalization and activation.
 
     Attributes:
         conv (nn.Conv2d): Convolutional layer.
@@ -48,8 +49,7 @@ class Conv(nn.Module):
     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.
+        """Initialize Conv layer with given parameters.
 
         Args:
             c1 (int): Number of input channels.
@@ -67,8 +67,7 @@ class Conv(nn.Module):
         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.
+        """Apply convolution, batch normalization and activation to input tensor.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -79,8 +78,7 @@ class Conv(nn.Module):
         return self.act(self.bn(self.conv(x)))
 
     def forward_fuse(self, x):
-        """
-        Apply convolution and activation without batch normalization.
+        """Apply convolution and activation without batch normalization.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -92,8 +90,7 @@ class Conv(nn.Module):
 
 
 class Conv2(Conv):
-    """
-    Simplified RepConv module with Conv fusing.
+    """Simplified RepConv module with Conv fusing.
 
     Attributes:
         conv (nn.Conv2d): Main 3x3 convolutional layer.
@@ -103,8 +100,7 @@ class Conv2(Conv):
     """
 
     def __init__(self, c1, c2, k=3, s=1, p=None, g=1, d=1, act=True):
-        """
-        Initialize Conv2 layer with given parameters.
+        """Initialize Conv2 layer with given parameters.
 
         Args:
             c1 (int): Number of input channels.
@@ -120,8 +116,7 @@ class Conv2(Conv):
         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.
+        """Apply convolution, batch normalization and activation to input tensor.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -132,8 +127,7 @@ class Conv2(Conv):
         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.
+        """Apply fused convolution, batch normalization and activation to input tensor.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -154,8 +148,7 @@ class Conv2(Conv):
 
 
 class LightConv(nn.Module):
-    """
-    Light convolution module with 1x1 and depthwise convolutions.
+    """Light convolution module with 1x1 and depthwise convolutions.
 
     This implementation is based on the PaddleDetection HGNetV2 backbone.
 
@@ -165,8 +158,7 @@ class LightConv(nn.Module):
     """
 
     def __init__(self, c1, c2, k=1, act=nn.ReLU()):
-        """
-        Initialize LightConv layer with given parameters.
+        """Initialize LightConv layer with given parameters.
 
         Args:
             c1 (int): Number of input channels.
@@ -179,8 +171,7 @@ class LightConv(nn.Module):
         self.conv2 = DWConv(c2, c2, k, act=act)
 
     def forward(self, x):
-        """
-        Apply 2 convolutions to input tensor.
+        """Apply 2 convolutions to input tensor.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -195,8 +186,7 @@ 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.
+        """Initialize depth-wise convolution with given parameters.
 
         Args:
             c1 (int): Number of input channels.
@@ -213,8 +203,7 @@ 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.
+        """Initialize depth-wise transpose convolution with given parameters.
 
         Args:
             c1 (int): Number of input channels.
@@ -228,8 +217,7 @@ class DWConvTranspose2d(nn.ConvTranspose2d):
 
 
 class ConvTranspose(nn.Module):
-    """
-    Convolution transpose module with optional batch normalization and activation.
+    """Convolution transpose module with optional batch normalization and activation.
 
     Attributes:
         conv_transpose (nn.ConvTranspose2d): Transposed convolution layer.
@@ -241,8 +229,7 @@ class ConvTranspose(nn.Module):
     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.
+        """Initialize ConvTranspose layer with given parameters.
 
         Args:
             c1 (int): Number of input channels.
@@ -259,8 +246,7 @@ class ConvTranspose(nn.Module):
         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.
+        """Apply transposed convolution, batch normalization and activation to input.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -271,8 +257,7 @@ class ConvTranspose(nn.Module):
         return self.act(self.bn(self.conv_transpose(x)))
 
     def forward_fuse(self, x):
-        """
-        Apply activation and convolution transpose operation to input.
+        """Apply activation and convolution transpose operation to input.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -284,8 +269,7 @@ class ConvTranspose(nn.Module):
 
 
 class Focus(nn.Module):
-    """
-    Focus module for concentrating feature information.
+    """Focus module for concentrating feature information.
 
     Slices input tensor into 4 parts and concatenates them in the channel dimension.
 
@@ -294,8 +278,7 @@ class Focus(nn.Module):
     """
 
     def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
-        """
-        Initialize Focus module with given parameters.
+        """Initialize Focus module with given parameters.
 
         Args:
             c1 (int): Number of input channels.
@@ -311,10 +294,9 @@ class Focus(nn.Module):
         # self.contract = Contract(gain=2)
 
     def forward(self, x):
-        """
-        Apply Focus operation and convolution to input tensor.
+        """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).
+        Input shape is (B, C, W, H) and output shape is (B, 4C, W/2, H/2).
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -327,8 +309,7 @@ class Focus(nn.Module):
 
 
 class GhostConv(nn.Module):
-    """
-    Ghost Convolution module.
+    """Ghost Convolution module.
 
     Generates more features with fewer parameters by using cheap operations.
 
@@ -341,8 +322,7 @@ class GhostConv(nn.Module):
     """
 
     def __init__(self, c1, c2, k=1, s=1, g=1, act=True):
-        """
-        Initialize Ghost Convolution module with given parameters.
+        """Initialize Ghost Convolution module with given parameters.
 
         Args:
             c1 (int): Number of input channels.
@@ -358,8 +338,7 @@ class GhostConv(nn.Module):
         self.cv2 = Conv(c_, c_, 5, 1, None, c_, act=act)
 
     def forward(self, x):
-        """
-        Apply Ghost Convolution to input tensor.
+        """Apply Ghost Convolution to input tensor.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -372,8 +351,7 @@ class GhostConv(nn.Module):
 
 
 class RepConv(nn.Module):
-    """
-    RepConv module with training and deploy modes.
+    """RepConv module with training and deploy modes.
 
     This module is used in RT-DETR and can fuse convolutions during inference for efficiency.
 
@@ -391,8 +369,7 @@ class RepConv(nn.Module):
     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.
+        """Initialize RepConv module with given parameters.
 
         Args:
             c1 (int): Number of input channels.
@@ -418,8 +395,7 @@ class RepConv(nn.Module):
         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.
+        """Forward pass for deploy mode.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -430,8 +406,7 @@ class RepConv(nn.Module):
         return self.act(self.conv(x))
 
     def forward(self, x):
-        """
-        Forward pass for training mode.
+        """Forward pass for training mode.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -443,8 +418,7 @@ class RepConv(nn.Module):
         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.
+        """Calculate equivalent kernel and bias by fusing convolutions.
 
         Returns:
             (torch.Tensor): Equivalent kernel
@@ -457,8 +431,7 @@ class RepConv(nn.Module):
 
     @staticmethod
     def _pad_1x1_to_3x3_tensor(kernel1x1):
-        """
-        Pad a 1x1 kernel to 3x3 size.
+        """Pad a 1x1 kernel to 3x3 size.
 
         Args:
             kernel1x1 (torch.Tensor): 1x1 convolution kernel.
@@ -472,15 +445,14 @@ class RepConv(nn.Module):
             return torch.nn.functional.pad(kernel1x1, [1, 1, 1, 1])
 
     def _fuse_bn_tensor(self, branch):
-        """
-        Fuse batch normalization with convolution weights.
+        """Fuse batch normalization with convolution weights.
 
         Args:
             branch (Conv | nn.BatchNorm2d | None): Branch to fuse.
 
         Returns:
-            (torch.Tensor): Fused kernel
-            (torch.Tensor): Fused bias
+            kernel (torch.Tensor): Fused kernel.
+            bias (torch.Tensor): Fused bias.
         """
         if branch is None:
             return 0, 0
@@ -538,8 +510,7 @@ class RepConv(nn.Module):
 
 
 class ChannelAttention(nn.Module):
-    """
-    Channel-attention module for feature recalibration.
+    """Channel-attention module for feature recalibration.
 
     Applies attention weights to channels based on global average pooling.
 
@@ -553,8 +524,7 @@ class ChannelAttention(nn.Module):
     """
 
     def __init__(self, channels: int) -> None:
-        """
-        Initialize Channel-attention module.
+        """Initialize Channel-attention module.
 
         Args:
             channels (int): Number of input channels.
@@ -565,8 +535,7 @@ class ChannelAttention(nn.Module):
         self.act = nn.Sigmoid()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Apply channel attention to input tensor.
+        """Apply channel attention to input tensor.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -578,8 +547,7 @@ class ChannelAttention(nn.Module):
 
 
 class SpatialAttention(nn.Module):
-    """
-    Spatial-attention module for feature recalibration.
+    """Spatial-attention module for feature recalibration.
 
     Applies attention weights to spatial dimensions based on channel statistics.
 
@@ -589,8 +557,7 @@ class SpatialAttention(nn.Module):
     """
 
     def __init__(self, kernel_size=7):
-        """
-        Initialize Spatial-attention module.
+        """Initialize Spatial-attention module.
 
         Args:
             kernel_size (int): Size of the convolutional kernel (3 or 7).
@@ -602,8 +569,7 @@ class SpatialAttention(nn.Module):
         self.act = nn.Sigmoid()
 
     def forward(self, x):
-        """
-        Apply spatial attention to input tensor.
+        """Apply spatial attention to input tensor.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -615,8 +581,7 @@ class SpatialAttention(nn.Module):
 
 
 class CBAM(nn.Module):
-    """
-    Convolutional Block Attention Module.
+    """Convolutional Block Attention Module.
 
     Combines channel and spatial attention mechanisms for comprehensive feature refinement.
 
@@ -626,8 +591,7 @@ class CBAM(nn.Module):
     """
 
     def __init__(self, c1, kernel_size=7):
-        """
-        Initialize CBAM with given parameters.
+        """Initialize CBAM with given parameters.
 
         Args:
             c1 (int): Number of input channels.
@@ -638,8 +602,7 @@ class CBAM(nn.Module):
         self.spatial_attention = SpatialAttention(kernel_size)
 
     def forward(self, x):
-        """
-        Apply channel and spatial attention sequentially to input tensor.
+        """Apply channel and spatial attention sequentially to input tensor.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -651,16 +614,14 @@ class CBAM(nn.Module):
 
 
 class Concat(nn.Module):
-    """
-    Concatenate a list of tensors along specified dimension.
+    """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.
+        """Initialize Concat module.
 
         Args:
             dimension (int): Dimension along which to concatenate tensors.
@@ -668,12 +629,11 @@ class Concat(nn.Module):
         super().__init__()
         self.d = dimension
 
-    def forward(self, x):
-        """
-        Concatenate input tensors along specified dimension.
+    def forward(self, x: list[torch.Tensor]):
+        """Concatenate input tensors along specified dimension.
 
         Args:
-            x (List[torch.Tensor]): List of input tensors.
+            x (list[torch.Tensor]): List of input tensors.
 
         Returns:
             (torch.Tensor): Concatenated tensor.
@@ -682,16 +642,14 @@ class Concat(nn.Module):
 
 
 class Index(nn.Module):
-    """
-    Returns a particular index of the input.
+    """Returns a particular index of the input.
 
     Attributes:
         index (int): Index to select from input.
     """
 
     def __init__(self, index=0):
-        """
-        Initialize Index module.
+        """Initialize Index module.
 
         Args:
             index (int): Index to select from input.
@@ -699,12 +657,11 @@ class Index(nn.Module):
         super().__init__()
         self.index = index
 
-    def forward(self, x):
-        """
-        Select and return a particular index from input.
+    def forward(self, x: list[torch.Tensor]):
+        """Select and return a particular index from input.
 
         Args:
-            x (List[torch.Tensor]): List of input tensors.
+            x (list[torch.Tensor]): List of input tensors.
 
         Returns:
             (torch.Tensor): Selected tensor.