ultralytics 8.3.89__py3-none-any.whl → 8.3.91__py3-none-any.whl

ultralytics/nn/modules/conv.py

@@ -35,40 +35,112 @@ def autopad(k, p=None, d=1):  # kernel, padding, dilation
 
 
 class Conv(nn.Module):
-    """Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)."""
+    """
+    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 arguments including activation."""
+        """
+        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."""
+        """
+        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."""
+        """
+        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."""
+    """
+    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 Conv layer with given arguments including activation."""
+        """
+        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."""
+        """
+        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."""
+        """
+        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):
@@ -83,106 +155,257 @@ class Conv2(Conv):
 
 class LightConv(nn.Module):
     """
-    Light convolution with args(ch_in, ch_out, kernel).
+    Light convolution module with 1x1 and depthwise convolutions.
+
+    This implementation is based on the PaddleDetection HGNetV2 backbone.
 
-    https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/backbones/hgnet_v2.py
+    Attributes:
+        conv1 (Conv): 1x1 convolution layer.
+        conv2 (DWConv): Depthwise convolution layer.
     """
 
     def __init__(self, c1, c2, k=1, act=nn.ReLU()):
-        """Initialize Conv layer with given arguments including activation."""
+        """
+        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."""
+        """
+        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."""
+    """Depth-wise convolution module."""
 
-    def __init__(self, c1, c2, k=1, s=1, d=1, act=True):  # ch_in, ch_out, kernel, stride, dilation, activation
-        """Initialize Depth-wise convolution with given parameters."""
+    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."""
+    """Depth-wise transpose convolution module."""
 
-    def __init__(self, c1, c2, k=1, s=1, p1=0, p2=0):  # ch_in, ch_out, kernel, stride, padding, padding_out
-        """Initialize DWConvTranspose2d class with given parameters."""
+    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 2d layer."""
+    """
+    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 ConvTranspose2d layer with batch normalization and activation function."""
+        """
+        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):
-        """Applies transposed convolutions, batch normalization and activation to input."""
+        """
+        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):
-        """Applies activation and convolution transpose operation to input."""
+        """
+        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 wh information into c-space."""
+    """
+    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):
-        """Initializes Focus object with user defined channel, convolution, padding, group and activation values."""
+        """
+        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):
         """
-        Applies convolution to concatenated tensor and returns the output.
+        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 https://github.com/huawei-noah/ghostnet."""
+    """
+    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/ghostnet
+    """
 
     def __init__(self, c1, c2, k=1, s=1, g=1, act=True):
-        """Initializes Ghost Convolution module with primary and cheap operations for efficient feature learning."""
+        """
+        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):
-        """Forward propagation through a Ghost Bottleneck layer with skip connection."""
+        """
+        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 is a basic rep-style block, including training and deploy status.
+    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).
 
-    This module is used in RT-DETR.
-    Based on https://github.com/DingXiaoH/RepVGG/blob/main/repvgg.py
+    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):
-        """Initializes Light Convolution layer with inputs, outputs & optional activation function."""
+        """
+        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
@@ -195,16 +418,39 @@ 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 process."""
+        """
+        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 process."""
+        """
+        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):
-        """Returns equivalent kernel and bias by adding 3x3 kernel, 1x1 kernel and identity kernel with their biases."""
+        """
+        Calculate equivalent kernel and bias by fusing convolutions.
+
+        Returns:
+            (tuple): Tuple containing:
+                - Equivalent kernel (torch.Tensor)
+                - Equivalent bias (torch.Tensor)
+        """
         kernel3x3, bias3x3 = self._fuse_bn_tensor(self.conv1)
         kernel1x1, bias1x1 = self._fuse_bn_tensor(self.conv2)
         kernelid, biasid = self._fuse_bn_tensor(self.bn)
@@ -212,14 +458,32 @@ class RepConv(nn.Module):
 
     @staticmethod
     def _pad_1x1_to_3x3_tensor(kernel1x1):
-        """Pads a 1x1 tensor to a 3x3 tensor."""
+        """
+        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):
-        """Generates appropriate kernels and biases for convolution by fusing branches of the neural network."""
+        """
+        Fuse batch normalization with convolution weights.
+
+        Args:
+            branch (Conv | nn.BatchNorm2d | None): Branch to fuse.
+
+        Returns:
+            (tuple): Tuple containing:
+                - Fused kernel (torch.Tensor)
+                - Fused bias (torch.Tensor)
+        """
         if branch is None:
             return 0, 0
         if isinstance(branch, Conv):
@@ -247,7 +511,7 @@ class RepConv(nn.Module):
         return kernel * t, beta - running_mean * gamma / std
 
     def fuse_convs(self):
-        """Combines two convolution layers into a single layer and removes unused attributes from the class."""
+        """Fuse convolutions for inference by creating a single equivalent convolution."""
         if hasattr(self, "conv"):
             return
         kernel, bias = self.get_equivalent_kernel_bias()
@@ -276,25 +540,63 @@ class RepConv(nn.Module):
 
 
 class ChannelAttention(nn.Module):
-    """Channel-attention module https://github.com/open-mmlab/mmdetection/tree/v3.0.0rc1/configs/rtmdet."""
+    """
+    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:
-        """Initializes the class and sets the basic configurations and instance variables required."""
+        """
+        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:
-        """Applies forward pass using activation on convolutions of the input, optionally using batch normalization."""
+        """
+        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."""
+    """
+    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 with kernel size argument."""
+        """
+        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
@@ -302,49 +604,111 @@ class SpatialAttention(nn.Module):
         self.act = nn.Sigmoid()
 
     def forward(self, x):
-        """Apply channel and spatial attention on input for feature recalibration."""
+        """
+        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."""
+    """
+    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 input channel (c1) and kernel size."""
+        """
+        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):
-        """Applies the forward pass through C1 module."""
+        """
+        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 dimension."""
+    """
+    Concatenate a list of tensors along specified dimension.
+
+    Attributes:
+        d (int): Dimension along which to concatenate tensors.
+    """
 
     def __init__(self, dimension=1):
-        """Concatenates a list of tensors along a specified dimension."""
+        """
+        Initialize Concat module.
+
+        Args:
+            dimension (int): Dimension along which to concatenate tensors.
+        """
         super().__init__()
         self.d = dimension
 
     def forward(self, x):
-        """Forward pass for the YOLOv8 mask Proto module."""
+        """
+        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."""
+    """
+    Returns a particular index of the input.
+
+    Attributes:
+        index (int): Index to select from input.
+    """
 
     def __init__(self, index=0):
-        """Returns a particular index of the input."""
+        """
+        Initialize Index module.
+
+        Args:
+            index (int): Index to select from input.
+        """
         super().__init__()
         self.index = index
 
     def forward(self, x):
         """
-        Forward pass.
+        Select and return a particular index from input.
+
+        Args:
+            x (List[torch.Tensor]): List of input tensors.
 
-        Expects a list of tensors as input.
+        Returns:
+            (torch.Tensor): Selected tensor.
         """
         return x[self.index]