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

ultralytics/nn/modules/block.py

@@ -69,7 +69,7 @@ class DFL(nn.Module):
         self.c1 = c1
 
     def forward(self, x):
-        """Applies a transformer layer on input tensor 'x' and returns a tensor."""
+        """Apply the DFL module to input tensor and return transformed output."""
         b, _, a = x.shape  # batch, channels, anchors
         return self.conv(x.view(b, 4, self.c1, a).transpose(2, 1).softmax(1)).view(b, 4, a)
         # return self.conv(x.view(b, self.c1, 4, a).softmax(1)).view(b, 4, a)
@@ -80,9 +80,12 @@ class Proto(nn.Module):
 
     def __init__(self, c1, c_=256, c2=32):
         """
-        Initializes the YOLOv8 mask Proto module with specified number of protos and masks.
+        Initialize the YOLOv8 mask Proto module with specified number of protos and masks.
 
-        Input arguments are ch_in, number of protos, number of masks.
+        Args:
+            c1 (int): Input channels.
+            c_ (int): Intermediate channels.
+            c2 (int): Output channels (number of protos).
         """
         super().__init__()
         self.cv1 = Conv(c1, c_, k=3)
@@ -91,7 +94,7 @@ class Proto(nn.Module):
         self.cv3 = Conv(c_, c2)
 
     def forward(self, x):
-        """Performs a forward pass through layers using an upsampled input image."""
+        """Perform a forward pass through layers using an upsampled input image."""
         return self.cv3(self.cv2(self.upsample(self.cv1(x))))
 
 
@@ -103,7 +106,14 @@ class HGStem(nn.Module):
     """
 
     def __init__(self, c1, cm, c2):
-        """Initialize the SPP layer with input/output channels and specified kernel sizes for max pooling."""
+        """
+        Initialize the StemBlock of PPHGNetV2.
+
+        Args:
+            c1 (int): Input channels.
+            cm (int): Middle channels.
+            c2 (int): Output channels.
+        """
         super().__init__()
         self.stem1 = Conv(c1, cm, 3, 2, act=nn.ReLU())
         self.stem2a = Conv(cm, cm // 2, 2, 1, 0, act=nn.ReLU())
@@ -134,7 +144,19 @@ class HGBlock(nn.Module):
     """
 
     def __init__(self, c1, cm, c2, k=3, n=6, lightconv=False, shortcut=False, act=nn.ReLU()):
-        """Initializes a CSP Bottleneck with 1 convolution using specified input and output channels."""
+        """
+        Initialize HGBlock with specified parameters.
+
+        Args:
+            c1 (int): Input channels.
+            cm (int): Middle channels.
+            c2 (int): Output channels.
+            k (int): Kernel size.
+            n (int): Number of LightConv or Conv blocks.
+            lightconv (bool): Whether to use LightConv.
+            shortcut (bool): Whether to use shortcut connection.
+            act (nn.Module): Activation function.
+        """
         super().__init__()
         block = LightConv if lightconv else Conv
         self.m = nn.ModuleList(block(c1 if i == 0 else cm, cm, k=k, act=act) for i in range(n))
@@ -154,7 +176,14 @@ class SPP(nn.Module):
     """Spatial Pyramid Pooling (SPP) layer https://arxiv.org/abs/1406.4729."""
 
     def __init__(self, c1, c2, k=(5, 9, 13)):
-        """Initialize the SPP layer with input/output channels and pooling kernel sizes."""
+        """
+        Initialize the SPP layer with input/output channels and pooling kernel sizes.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            k (Tuple[int, int, int]): Kernel sizes for max pooling.
+        """
         super().__init__()
         c_ = c1 // 2  # hidden channels
         self.cv1 = Conv(c1, c_, 1, 1)
@@ -172,9 +201,15 @@ class SPPF(nn.Module):
 
     def __init__(self, c1, c2, k=5):
         """
-        Initializes the SPPF layer with given input/output channels and kernel size.
+        Initialize the SPPF layer with given input/output channels and kernel size.
 
-        This module is equivalent to SPP(k=(5, 9, 13)).
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            k (int): Kernel size.
+
+        Notes:
+            This module is equivalent to SPP(k=(5, 9, 13)).
         """
         super().__init__()
         c_ = c1 // 2  # hidden channels
@@ -183,7 +218,7 @@ class SPPF(nn.Module):
         self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
 
     def forward(self, x):
-        """Forward pass through Ghost Convolution block."""
+        """Apply sequential pooling operations to input and return concatenated feature maps."""
         y = [self.cv1(x)]
         y.extend(self.m(y[-1]) for _ in range(3))
         return self.cv2(torch.cat(y, 1))
@@ -193,13 +228,20 @@ class C1(nn.Module):
     """CSP Bottleneck with 1 convolution."""
 
     def __init__(self, c1, c2, n=1):
-        """Initializes the CSP Bottleneck with configurations for 1 convolution with arguments ch_in, ch_out, number."""
+        """
+        Initialize the CSP Bottleneck with 1 convolution.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            n (int): Number of convolutions.
+        """
         super().__init__()
         self.cv1 = Conv(c1, c2, 1, 1)
         self.m = nn.Sequential(*(Conv(c2, c2, 3) for _ in range(n)))
 
     def forward(self, x):
-        """Applies cross-convolutions to input in the C3 module."""
+        """Apply convolution and residual connection to input tensor."""
         y = self.cv1(x)
         return self.m(y) + y
 
@@ -208,7 +250,17 @@ class C2(nn.Module):
     """CSP Bottleneck with 2 convolutions."""
 
     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
-        """Initializes a CSP Bottleneck with 2 convolutions and optional shortcut connection."""
+        """
+        Initialize a CSP Bottleneck with 2 convolutions.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            n (int): Number of Bottleneck blocks.
+            shortcut (bool): Whether to use shortcut connections.
+            g (int): Groups for convolutions.
+            e (float): Expansion ratio.
+        """
         super().__init__()
         self.c = int(c2 * e)  # hidden channels
         self.cv1 = Conv(c1, 2 * self.c, 1, 1)
@@ -226,7 +278,17 @@ class C2f(nn.Module):
     """Faster Implementation of CSP Bottleneck with 2 convolutions."""
 
     def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
-        """Initializes a CSP bottleneck with 2 convolutions and n Bottleneck blocks for faster processing."""
+        """
+        Initialize a CSP bottleneck with 2 convolutions.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            n (int): Number of Bottleneck blocks.
+            shortcut (bool): Whether to use shortcut connections.
+            g (int): Groups for convolutions.
+            e (float): Expansion ratio.
+        """
         super().__init__()
         self.c = int(c2 * e)  # hidden channels
         self.cv1 = Conv(c1, 2 * self.c, 1, 1)
@@ -251,7 +313,17 @@ class C3(nn.Module):
     """CSP Bottleneck with 3 convolutions."""
 
     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
-        """Initialize the CSP Bottleneck with given channels, number, shortcut, groups, and expansion values."""
+        """
+        Initialize the CSP Bottleneck with 3 convolutions.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            n (int): Number of Bottleneck blocks.
+            shortcut (bool): Whether to use shortcut connections.
+            g (int): Groups for convolutions.
+            e (float): Expansion ratio.
+        """
         super().__init__()
         c_ = int(c2 * e)  # hidden channels
         self.cv1 = Conv(c1, c_, 1, 1)
@@ -260,7 +332,7 @@ class C3(nn.Module):
         self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, k=((1, 1), (3, 3)), e=1.0) for _ in range(n)))
 
     def forward(self, x):
-        """Forward pass through the CSP bottleneck with 2 convolutions."""
+        """Forward pass through the CSP bottleneck with 3 convolutions."""
         return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))
 
 
@@ -268,7 +340,17 @@ class C3x(C3):
     """C3 module with cross-convolutions."""
 
     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
-        """Initialize C3TR instance and set default parameters."""
+        """
+        Initialize C3 module with cross-convolutions.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            n (int): Number of Bottleneck blocks.
+            shortcut (bool): Whether to use shortcut connections.
+            g (int): Groups for convolutions.
+            e (float): Expansion ratio.
+        """
         super().__init__(c1, c2, n, shortcut, g, e)
         self.c_ = int(c2 * e)
         self.m = nn.Sequential(*(Bottleneck(self.c_, self.c_, shortcut, g, k=((1, 3), (3, 1)), e=1) for _ in range(n)))
@@ -278,7 +360,15 @@ class RepC3(nn.Module):
     """Rep C3."""
 
     def __init__(self, c1, c2, n=3, e=1.0):
-        """Initialize CSP Bottleneck with a single convolution using input channels, output channels, and number."""
+        """
+        Initialize CSP Bottleneck with a single convolution.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            n (int): Number of RepConv blocks.
+            e (float): Expansion ratio.
+        """
         super().__init__()
         c_ = int(c2 * e)  # hidden channels
         self.cv1 = Conv(c1, c_, 1, 1)
@@ -287,7 +377,7 @@ class RepC3(nn.Module):
         self.cv3 = Conv(c_, c2, 1, 1) if c_ != c2 else nn.Identity()
 
     def forward(self, x):
-        """Forward pass of RT-DETR neck layer."""
+        """Forward pass of RepC3 module."""
         return self.cv3(self.m(self.cv1(x)) + self.cv2(x))
 
 
@@ -295,7 +385,17 @@ class C3TR(C3):
     """C3 module with TransformerBlock()."""
 
     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
-        """Initialize C3Ghost module with GhostBottleneck()."""
+        """
+        Initialize C3 module with TransformerBlock.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            n (int): Number of Transformer blocks.
+            shortcut (bool): Whether to use shortcut connections.
+            g (int): Groups for convolutions.
+            e (float): Expansion ratio.
+        """
         super().__init__(c1, c2, n, shortcut, g, e)
         c_ = int(c2 * e)
         self.m = TransformerBlock(c_, c_, 4, n)
@@ -305,7 +405,17 @@ class C3Ghost(C3):
     """C3 module with GhostBottleneck()."""
 
     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
-        """Initialize 'SPP' module with various pooling sizes for spatial pyramid pooling."""
+        """
+        Initialize C3 module with GhostBottleneck.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            n (int): Number of Ghost bottleneck blocks.
+            shortcut (bool): Whether to use shortcut connections.
+            g (int): Groups for convolutions.
+            e (float): Expansion ratio.
+        """
         super().__init__(c1, c2, n, shortcut, g, e)
         c_ = int(c2 * e)  # hidden channels
         self.m = nn.Sequential(*(GhostBottleneck(c_, c_) for _ in range(n)))
@@ -315,7 +425,15 @@ class GhostBottleneck(nn.Module):
     """Ghost Bottleneck https://github.com/huawei-noah/ghostnet."""
 
     def __init__(self, c1, c2, k=3, s=1):
-        """Initializes GhostBottleneck module with arguments ch_in, ch_out, kernel, stride."""
+        """
+        Initialize Ghost Bottleneck module.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            k (int): Kernel size.
+            s (int): Stride.
+        """
         super().__init__()
         c_ = c2 // 2
         self.conv = nn.Sequential(
@@ -328,7 +446,7 @@ class GhostBottleneck(nn.Module):
         )
 
     def forward(self, x):
-        """Applies skip connection and concatenation to input tensor."""
+        """Apply skip connection and concatenation to input tensor."""
         return self.conv(x) + self.shortcut(x)
 
 
@@ -336,7 +454,17 @@ class Bottleneck(nn.Module):
     """Standard bottleneck."""
 
     def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):
-        """Initializes a standard bottleneck module with optional shortcut connection and configurable parameters."""
+        """
+        Initialize a standard bottleneck module.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            shortcut (bool): Whether to use shortcut connection.
+            g (int): Groups for convolutions.
+            k (Tuple[int, int]): Kernel sizes for convolutions.
+            e (float): Expansion ratio.
+        """
         super().__init__()
         c_ = int(c2 * e)  # hidden channels
         self.cv1 = Conv(c1, c_, k[0], 1)
@@ -344,7 +472,7 @@ class Bottleneck(nn.Module):
         self.add = shortcut and c1 == c2
 
     def forward(self, x):
-        """Applies the YOLO FPN to input data."""
+        """Apply bottleneck with optional shortcut connection."""
         return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
 
 
@@ -352,7 +480,17 @@ class BottleneckCSP(nn.Module):
     """CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks."""
 
     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
-        """Initializes the CSP Bottleneck given arguments for ch_in, ch_out, number, shortcut, groups, expansion."""
+        """
+        Initialize CSP Bottleneck.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            n (int): Number of Bottleneck blocks.
+            shortcut (bool): Whether to use shortcut connections.
+            g (int): Groups for convolutions.
+            e (float): Expansion ratio.
+        """
         super().__init__()
         c_ = int(c2 * e)  # hidden channels
         self.cv1 = Conv(c1, c_, 1, 1)
@@ -364,7 +502,7 @@ class BottleneckCSP(nn.Module):
         self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
 
     def forward(self, x):
-        """Applies a CSP bottleneck with 3 convolutions."""
+        """Apply CSP bottleneck with 3 convolutions."""
         y1 = self.cv3(self.m(self.cv1(x)))
         y2 = self.cv2(x)
         return self.cv4(self.act(self.bn(torch.cat((y1, y2), 1))))
@@ -374,7 +512,15 @@ class ResNetBlock(nn.Module):
     """ResNet block with standard convolution layers."""
 
     def __init__(self, c1, c2, s=1, e=4):
-        """Initialize convolution with given parameters."""
+        """
+        Initialize ResNet block.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            s (int): Stride.
+            e (int): Expansion ratio.
+        """
         super().__init__()
         c3 = e * c2
         self.cv1 = Conv(c1, c2, k=1, s=1, act=True)
@@ -391,7 +537,17 @@ class ResNetLayer(nn.Module):
     """ResNet layer with multiple ResNet blocks."""
 
     def __init__(self, c1, c2, s=1, is_first=False, n=1, e=4):
-        """Initializes the ResNetLayer given arguments."""
+        """
+        Initialize ResNet layer.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            s (int): Stride.
+            is_first (bool): Whether this is the first layer.
+            n (int): Number of ResNet blocks.
+            e (int): Expansion ratio.
+        """
         super().__init__()
         self.is_first = is_first
 
@@ -413,7 +569,17 @@ class MaxSigmoidAttnBlock(nn.Module):
     """Max Sigmoid attention block."""
 
     def __init__(self, c1, c2, nh=1, ec=128, gc=512, scale=False):
-        """Initializes MaxSigmoidAttnBlock with specified arguments."""
+        """
+        Initialize MaxSigmoidAttnBlock.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            nh (int): Number of heads.
+            ec (int): Embedding channels.
+            gc (int): Guide channels.
+            scale (bool): Whether to use learnable scale parameter.
+        """
         super().__init__()
         self.nh = nh
         self.hc = c2 // nh
@@ -424,7 +590,16 @@ class MaxSigmoidAttnBlock(nn.Module):
         self.scale = nn.Parameter(torch.ones(1, nh, 1, 1)) if scale else 1.0
 
     def forward(self, x, guide):
-        """Forward process."""
+        """
+        Forward pass of MaxSigmoidAttnBlock.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+            guide (torch.Tensor): Guide tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor after attention.
+        """
         bs, _, h, w = x.shape
 
         guide = self.gl(guide)
@@ -448,7 +623,20 @@ class C2fAttn(nn.Module):
     """C2f module with an additional attn module."""
 
     def __init__(self, c1, c2, n=1, ec=128, nh=1, gc=512, shortcut=False, g=1, e=0.5):
-        """Initializes C2f module with attention mechanism for enhanced feature extraction and processing."""
+        """
+        Initialize C2f module with attention mechanism.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            n (int): Number of Bottleneck blocks.
+            ec (int): Embedding channels for attention.
+            nh (int): Number of heads for attention.
+            gc (int): Guide channels for attention.
+            shortcut (bool): Whether to use shortcut connections.
+            g (int): Groups for convolutions.
+            e (float): Expansion ratio.
+        """
         super().__init__()
         self.c = int(c2 * e)  # hidden channels
         self.cv1 = Conv(c1, 2 * self.c, 1, 1)
@@ -457,14 +645,32 @@ class C2fAttn(nn.Module):
         self.attn = MaxSigmoidAttnBlock(self.c, self.c, gc=gc, ec=ec, nh=nh)
 
     def forward(self, x, guide):
-        """Forward pass through C2f layer."""
+        """
+        Forward pass through C2f layer with attention.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+            guide (torch.Tensor): Guide tensor for attention.
+
+        Returns:
+            (torch.Tensor): Output tensor after processing.
+        """
         y = list(self.cv1(x).chunk(2, 1))
         y.extend(m(y[-1]) for m in self.m)
         y.append(self.attn(y[-1], guide))
         return self.cv2(torch.cat(y, 1))
 
     def forward_split(self, x, guide):
-        """Forward pass using split() instead of chunk()."""
+        """
+        Forward pass using split() instead of chunk().
+
+        Args:
+            x (torch.Tensor): Input tensor.
+            guide (torch.Tensor): Guide tensor for attention.
+
+        Returns:
+            (torch.Tensor): Output tensor after processing.
+        """
         y = list(self.cv1(x).split((self.c, self.c), 1))
         y.extend(m(y[-1]) for m in self.m)
         y.append(self.attn(y[-1], guide))
@@ -475,7 +681,17 @@ class ImagePoolingAttn(nn.Module):
     """ImagePoolingAttn: Enhance the text embeddings with image-aware information."""
 
     def __init__(self, ec=256, ch=(), ct=512, nh=8, k=3, scale=False):
-        """Initializes ImagePoolingAttn with specified arguments."""
+        """
+        Initialize ImagePoolingAttn module.
+
+        Args:
+            ec (int): Embedding channels.
+            ch (Tuple): Channel dimensions for feature maps.
+            ct (int): Channel dimension for text embeddings.
+            nh (int): Number of attention heads.
+            k (int): Kernel size for pooling.
+            scale (bool): Whether to use learnable scale parameter.
+        """
         super().__init__()
 
         nf = len(ch)
@@ -493,7 +709,16 @@ class ImagePoolingAttn(nn.Module):
         self.k = k
 
     def forward(self, x, text):
-        """Executes attention mechanism on input tensor x and guide tensor."""
+        """
+        Forward pass of ImagePoolingAttn.
+
+        Args:
+            x (List[torch.Tensor]): List of input feature maps.
+            text (torch.Tensor): Text embeddings.
+
+        Returns:
+            (torch.Tensor): Enhanced text embeddings.
+        """
         bs = x[0].shape[0]
         assert len(x) == self.nf
         num_patches = self.k**2
@@ -521,14 +746,23 @@ class ContrastiveHead(nn.Module):
     """Implements contrastive learning head for region-text similarity in vision-language models."""
 
     def __init__(self):
-        """Initializes ContrastiveHead with specified region-text similarity parameters."""
+        """Initialize ContrastiveHead with region-text similarity parameters."""
         super().__init__()
         # NOTE: use -10.0 to keep the init cls loss consistency with other losses
         self.bias = nn.Parameter(torch.tensor([-10.0]))
         self.logit_scale = nn.Parameter(torch.ones([]) * torch.tensor(1 / 0.07).log())
 
     def forward(self, x, w):
-        """Forward function of contrastive learning."""
+        """
+        Forward function of contrastive learning.
+
+        Args:
+            x (torch.Tensor): Image features.
+            w (torch.Tensor): Text features.
+
+        Returns:
+            (torch.Tensor): Similarity scores.
+        """
         x = F.normalize(x, dim=1, p=2)
         w = F.normalize(w, dim=-1, p=2)
         x = torch.einsum("bchw,bkc->bkhw", x, w)
@@ -544,7 +778,12 @@ class BNContrastiveHead(nn.Module):
     """
 
     def __init__(self, embed_dims: int):
-        """Initialize ContrastiveHead with region-text similarity parameters."""
+        """
+        Initialize BNContrastiveHead.
+
+        Args:
+            embed_dims (int): Embedding dimensions for features.
+        """
         super().__init__()
         self.norm = nn.BatchNorm2d(embed_dims)
         # NOTE: use -10.0 to keep the init cls loss consistency with other losses
@@ -553,7 +792,16 @@ class BNContrastiveHead(nn.Module):
         self.logit_scale = nn.Parameter(-1.0 * torch.ones([]))
 
     def forward(self, x, w):
-        """Forward function of contrastive learning."""
+        """
+        Forward function of contrastive learning with batch normalization.
+
+        Args:
+            x (torch.Tensor): Image features.
+            w (torch.Tensor): Text features.
+
+        Returns:
+            (torch.Tensor): Similarity scores.
+        """
         x = self.norm(x)
         w = F.normalize(w, dim=-1, p=2)
         x = torch.einsum("bchw,bkc->bkhw", x, w)
@@ -564,7 +812,17 @@ class RepBottleneck(Bottleneck):
     """Rep bottleneck."""
 
     def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):
-        """Initializes a RepBottleneck module with customizable in/out channels, shortcuts, groups and expansion."""
+        """
+        Initialize RepBottleneck.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            shortcut (bool): Whether to use shortcut connection.
+            g (int): Groups for convolutions.
+            k (Tuple[int, int]): Kernel sizes for convolutions.
+            e (float): Expansion ratio.
+        """
         super().__init__(c1, c2, shortcut, g, k, e)
         c_ = int(c2 * e)  # hidden channels
         self.cv1 = RepConv(c1, c_, k[0], 1)
@@ -574,7 +832,17 @@ class RepCSP(C3):
     """Repeatable Cross Stage Partial Network (RepCSP) module for efficient feature extraction."""
 
     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
-        """Initializes RepCSP layer with given channels, repetitions, shortcut, groups and expansion ratio."""
+        """
+        Initialize RepCSP layer.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            n (int): Number of RepBottleneck blocks.
+            shortcut (bool): Whether to use shortcut connections.
+            g (int): Groups for convolutions.
+            e (float): Expansion ratio.
+        """
         super().__init__(c1, c2, n, shortcut, g, e)
         c_ = int(c2 * e)  # hidden channels
         self.m = nn.Sequential(*(RepBottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
@@ -584,7 +852,16 @@ class RepNCSPELAN4(nn.Module):
     """CSP-ELAN."""
 
     def __init__(self, c1, c2, c3, c4, n=1):
-        """Initializes CSP-ELAN layer with specified channel sizes, repetitions, and convolutions."""
+        """
+        Initialize CSP-ELAN layer.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            c3 (int): Intermediate channels.
+            c4 (int): Intermediate channels for RepCSP.
+            n (int): Number of RepCSP blocks.
+        """
         super().__init__()
         self.c = c3 // 2
         self.cv1 = Conv(c1, c3, 1, 1)
@@ -609,7 +886,15 @@ class ELAN1(RepNCSPELAN4):
     """ELAN1 module with 4 convolutions."""
 
     def __init__(self, c1, c2, c3, c4):
-        """Initializes ELAN1 layer with specified channel sizes."""
+        """
+        Initialize ELAN1 layer.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            c3 (int): Intermediate channels.
+            c4 (int): Intermediate channels for convolutions.
+        """
         super().__init__(c1, c2, c3, c4)
         self.c = c3 // 2
         self.cv1 = Conv(c1, c3, 1, 1)
@@ -622,7 +907,13 @@ class AConv(nn.Module):
     """AConv."""
 
     def __init__(self, c1, c2):
-        """Initializes AConv module with convolution layers."""
+        """
+        Initialize AConv module.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+        """
         super().__init__()
         self.cv1 = Conv(c1, c2, 3, 2, 1)
 
@@ -636,7 +927,13 @@ class ADown(nn.Module):
     """ADown."""
 
     def __init__(self, c1, c2):
-        """Initializes ADown module with convolution layers to downsample input from channels c1 to c2."""
+        """
+        Initialize ADown module.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+        """
         super().__init__()
         self.c = c2 // 2
         self.cv1 = Conv(c1 // 2, self.c, 3, 2, 1)
@@ -656,7 +953,15 @@ class SPPELAN(nn.Module):
     """SPP-ELAN."""
 
     def __init__(self, c1, c2, c3, k=5):
-        """Initializes SPP-ELAN block with convolution and max pooling layers for spatial pyramid pooling."""
+        """
+        Initialize SPP-ELAN block.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            c3 (int): Intermediate channels.
+            k (int): Kernel size for max pooling.
+        """
         super().__init__()
         self.c = c3
         self.cv1 = Conv(c1, c3, 1, 1)
@@ -676,7 +981,17 @@ class CBLinear(nn.Module):
     """CBLinear."""
 
     def __init__(self, c1, c2s, k=1, s=1, p=None, g=1):
-        """Initializes the CBLinear module, passing inputs unchanged."""
+        """
+        Initialize CBLinear module.
+
+        Args:
+            c1 (int): Input channels.
+            c2s (List[int]): List of output channel sizes.
+            k (int): Kernel size.
+            s (int): Stride.
+            p (int | None): Padding.
+            g (int): Groups.
+        """
         super().__init__()
         self.c2s = c2s
         self.conv = nn.Conv2d(c1, sum(c2s), k, s, autopad(k, p), groups=g, bias=True)
@@ -690,12 +1005,25 @@ class CBFuse(nn.Module):
     """CBFuse."""
 
     def __init__(self, idx):
-        """Initializes CBFuse module with layer index for selective feature fusion."""
+        """
+        Initialize CBFuse module.
+
+        Args:
+            idx (List[int]): Indices for feature selection.
+        """
         super().__init__()
         self.idx = idx
 
     def forward(self, xs):
-        """Forward pass through CBFuse layer."""
+        """
+        Forward pass through CBFuse layer.
+
+        Args:
+            xs (List[torch.Tensor]): List of input tensors.
+
+        Returns:
+            (torch.Tensor): Fused output tensor.
+        """
         target_size = xs[-1].shape[2:]
         res = [F.interpolate(x[self.idx[i]], size=target_size, mode="nearest") for i, x in enumerate(xs[:-1])]
         return torch.sum(torch.stack(res + xs[-1:]), dim=0)
@@ -705,8 +1033,16 @@ class C3f(nn.Module):
     """Faster Implementation of CSP Bottleneck with 2 convolutions."""
 
     def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
-        """Initialize CSP bottleneck layer with two convolutions with arguments ch_in, ch_out, number, shortcut, groups,
-        expansion.
+        """
+        Initialize CSP bottleneck layer with two convolutions.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            n (int): Number of Bottleneck blocks.
+            shortcut (bool): Whether to use shortcut connections.
+            g (int): Groups for convolutions.
+            e (float): Expansion ratio.
         """
         super().__init__()
         c_ = int(c2 * e)  # hidden channels
@@ -716,7 +1052,7 @@ class C3f(nn.Module):
         self.m = nn.ModuleList(Bottleneck(c_, c_, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))
 
     def forward(self, x):
-        """Forward pass through C2f layer."""
+        """Forward pass through C3f layer."""
         y = [self.cv2(x), self.cv1(x)]
         y.extend(m(y[-1]) for m in self.m)
         return self.cv3(torch.cat(y, 1))
@@ -726,7 +1062,18 @@ class C3k2(C2f):
     """Faster Implementation of CSP Bottleneck with 2 convolutions."""
 
     def __init__(self, c1, c2, n=1, c3k=False, e=0.5, g=1, shortcut=True):
-        """Initializes the C3k2 module, a faster CSP Bottleneck with 2 convolutions and optional C3k blocks."""
+        """
+        Initialize C3k2 module.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            n (int): Number of blocks.
+            c3k (bool): Whether to use C3k blocks.
+            e (float): Expansion ratio.
+            g (int): Groups for convolutions.
+            shortcut (bool): Whether to use shortcut connections.
+        """
         super().__init__(c1, c2, n, shortcut, g, e)
         self.m = nn.ModuleList(
             C3k(self.c, self.c, 2, shortcut, g) if c3k else Bottleneck(self.c, self.c, shortcut, g) for _ in range(n)
@@ -737,7 +1084,18 @@ class C3k(C3):
     """C3k is a CSP bottleneck module with customizable kernel sizes for feature extraction in neural networks."""
 
     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, k=3):
-        """Initializes the C3k module with specified channels, number of layers, and configurations."""
+        """
+        Initialize C3k module.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            n (int): Number of Bottleneck blocks.
+            shortcut (bool): Whether to use shortcut connections.
+            g (int): Groups for convolutions.
+            e (float): Expansion ratio.
+            k (int): Kernel size.
+        """
         super().__init__(c1, c2, n, shortcut, g, e)
         c_ = int(c2 * e)  # hidden channels
         # self.m = nn.Sequential(*(RepBottleneck(c_, c_, shortcut, g, k=(k, k), e=1.0) for _ in range(n)))
@@ -748,7 +1106,12 @@ class RepVGGDW(torch.nn.Module):
     """RepVGGDW is a class that represents a depth wise separable convolutional block in RepVGG architecture."""
 
     def __init__(self, ed) -> None:
-        """Initializes RepVGGDW with depthwise separable convolutional layers for efficient processing."""
+        """
+        Initialize RepVGGDW module.
+
+        Args:
+            ed (int): Input and output channels.
+        """
         super().__init__()
         self.conv = Conv(ed, ed, 7, 1, 3, g=ed, act=False)
         self.conv1 = Conv(ed, ed, 3, 1, 1, g=ed, act=False)
@@ -757,7 +1120,7 @@ class RepVGGDW(torch.nn.Module):
 
     def forward(self, x):
         """
-        Performs a forward pass of the RepVGGDW block.
+        Perform a forward pass of the RepVGGDW block.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -769,7 +1132,7 @@ class RepVGGDW(torch.nn.Module):
 
     def forward_fuse(self, x):
         """
-        Performs a forward pass of the RepVGGDW block without fusing the convolutions.
+        Perform a forward pass of the RepVGGDW block without fusing the convolutions.
 
         Args:
             x (torch.Tensor): Input tensor.
@@ -782,7 +1145,7 @@ class RepVGGDW(torch.nn.Module):
     @torch.no_grad()
     def fuse(self):
         """
-        Fuses the convolutional layers in the RepVGGDW block.
+        Fuse the convolutional layers in the RepVGGDW block.
 
         This method fuses the convolutional layers and updates the weights and biases accordingly.
         """
@@ -819,7 +1182,16 @@ class CIB(nn.Module):
     """
 
     def __init__(self, c1, c2, shortcut=True, e=0.5, lk=False):
-        """Initializes the custom model with optional shortcut, scaling factor, and RepVGGDW layer."""
+        """
+        Initialize the CIB module.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            shortcut (bool): Whether to use shortcut connection.
+            e (float): Expansion ratio.
+            lk (bool): Whether to use RepVGGDW.
+        """
         super().__init__()
         c_ = int(c2 * e)  # hidden channels
         self.cv1 = nn.Sequential(
@@ -860,7 +1232,18 @@ class C2fCIB(C2f):
     """
 
     def __init__(self, c1, c2, n=1, shortcut=False, lk=False, g=1, e=0.5):
-        """Initializes the module with specified parameters for channel, shortcut, local key, groups, and expansion."""
+        """
+        Initialize C2fCIB module.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            n (int): Number of CIB modules.
+            shortcut (bool): Whether to use shortcut connection.
+            lk (bool): Whether to use local key connection.
+            g (int): Groups for convolutions.
+            e (float): Expansion ratio.
+        """
         super().__init__(c1, c2, n, shortcut, g, e)
         self.m = nn.ModuleList(CIB(self.c, self.c, shortcut, e=1.0, lk=lk) for _ in range(n))
 
@@ -885,7 +1268,14 @@ class Attention(nn.Module):
     """
 
     def __init__(self, dim, num_heads=8, attn_ratio=0.5):
-        """Initializes multi-head attention module with query, key, and value convolutions and positional encoding."""
+        """
+        Initialize multi-head attention module.
+
+        Args:
+            dim (int): Input dimension.
+            num_heads (int): Number of attention heads.
+            attn_ratio (float): Attention ratio for key dimension.
+        """
         super().__init__()
         self.num_heads = num_heads
         self.head_dim = dim // num_heads
@@ -944,7 +1334,15 @@ class PSABlock(nn.Module):
     """
 
     def __init__(self, c, attn_ratio=0.5, num_heads=4, shortcut=True) -> None:
-        """Initializes the PSABlock with attention and feed-forward layers for enhanced feature extraction."""
+        """
+        Initialize the PSABlock.
+
+        Args:
+            c (int): Input and output channels.
+            attn_ratio (float): Attention ratio for key dimension.
+            num_heads (int): Number of attention heads.
+            shortcut (bool): Whether to use shortcut connections.
+        """
         super().__init__()
 
         self.attn = Attention(c, attn_ratio=attn_ratio, num_heads=num_heads)
@@ -952,7 +1350,15 @@ class PSABlock(nn.Module):
         self.add = shortcut
 
     def forward(self, x):
-        """Executes a forward pass through PSABlock, applying attention and feed-forward layers to the input tensor."""
+        """
+        Execute a forward pass through PSABlock.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor after attention and feed-forward processing.
+        """
         x = x + self.attn(x) if self.add else self.attn(x)
         x = x + self.ffn(x) if self.add else self.ffn(x)
         return x
@@ -983,7 +1389,14 @@ class PSA(nn.Module):
     """
 
     def __init__(self, c1, c2, e=0.5):
-        """Initializes the PSA module with input/output channels and attention mechanism for feature extraction."""
+        """
+        Initialize PSA module.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            e (float): Expansion ratio.
+        """
         super().__init__()
         assert c1 == c2
         self.c = int(c1 * e)
@@ -994,7 +1407,15 @@ class PSA(nn.Module):
         self.ffn = nn.Sequential(Conv(self.c, self.c * 2, 1), Conv(self.c * 2, self.c, 1, act=False))
 
     def forward(self, x):
-        """Executes forward pass in PSA module, applying attention and feed-forward layers to the input tensor."""
+        """
+        Execute forward pass in PSA module.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor after attention and feed-forward processing.
+        """
         a, b = self.cv1(x).split((self.c, self.c), dim=1)
         b = b + self.attn(b)
         b = b + self.ffn(b)
@@ -1027,7 +1448,15 @@ class C2PSA(nn.Module):
     """
 
     def __init__(self, c1, c2, n=1, e=0.5):
-        """Initializes the C2PSA module with specified input/output channels, number of layers, and expansion ratio."""
+        """
+        Initialize C2PSA module.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            n (int): Number of PSABlock modules.
+            e (float): Expansion ratio.
+        """
         super().__init__()
         assert c1 == c2
         self.c = int(c1 * e)
@@ -1037,7 +1466,15 @@ class C2PSA(nn.Module):
         self.m = nn.Sequential(*(PSABlock(self.c, attn_ratio=0.5, num_heads=self.c // 64) for _ in range(n)))
 
     def forward(self, x):
-        """Processes the input tensor 'x' through a series of PSA blocks and returns the transformed tensor."""
+        """
+        Process the input tensor through a series of PSA blocks.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor after processing.
+        """
         a, b = self.cv1(x).split((self.c, self.c), dim=1)
         b = self.m(b)
         return self.cv2(torch.cat((a, b), 1))
@@ -1069,7 +1506,15 @@ class C2fPSA(C2f):
     """
 
     def __init__(self, c1, c2, n=1, e=0.5):
-        """Initializes the C2fPSA module, a variant of C2f with PSA blocks for enhanced feature extraction."""
+        """
+        Initialize C2fPSA module.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            n (int): Number of PSABlock modules.
+            e (float): Expansion ratio.
+        """
         assert c1 == c2
         super().__init__(c1, c2, n=n, e=e)
         self.m = nn.ModuleList(PSABlock(self.c, attn_ratio=0.5, num_heads=self.c // 64) for _ in range(n))
@@ -1100,13 +1545,29 @@ class SCDown(nn.Module):
     """
 
     def __init__(self, c1, c2, k, s):
-        """Initializes the SCDown module with specified input/output channels, kernel size, and stride."""
+        """
+        Initialize SCDown module.
+
+        Args:
+            c1 (int): Input channels.
+            c2 (int): Output channels.
+            k (int): Kernel size.
+            s (int): Stride.
+        """
         super().__init__()
         self.cv1 = Conv(c1, c2, 1, 1)
         self.cv2 = Conv(c2, c2, k=k, s=s, g=c2, act=False)
 
     def forward(self, x):
-        """Applies convolution and downsampling to the input tensor in the SCDown module."""
+        """
+        Apply convolution and downsampling to the input tensor.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Downsampled output tensor.
+        """
         return self.cv2(self.cv1(x))
 
 
@@ -1128,7 +1589,16 @@ class TorchVision(nn.Module):
     """
 
     def __init__(self, model, weights="DEFAULT", unwrap=True, truncate=2, split=False):
-        """Load the model and weights from torchvision."""
+        """
+        Load the model and weights from torchvision.
+
+        Args:
+            model (str): Name of the torchvision model to load.
+            weights (str): Pre-trained weights to load.
+            unwrap (bool): Whether to unwrap the model.
+            truncate (int): Number of layers to truncate.
+            split (bool): Whether to split the output.
+        """
         import torchvision  # scope for faster 'import ultralytics'
 
         super().__init__()
@@ -1147,7 +1617,15 @@ class TorchVision(nn.Module):
             self.m.head = self.m.heads = nn.Identity()
 
     def forward(self, x):
-        """Forward pass through the model."""
+        """
+        Forward pass through the model.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor | List[torch.Tensor]): Output tensor or list of tensors.
+        """
         if self.split:
             y = [x]
             y.extend(m(y[-1]) for m in self.m)
@@ -1184,7 +1662,7 @@ class AAttn(nn.Module):
 
     def __init__(self, dim, num_heads, area=1):
         """
-        Initializes an Area-attention module for YOLO models.
+        Initialize an Area-attention module for YOLO models.
 
         Args:
             dim (int): Number of hidden channels.
@@ -1203,7 +1681,15 @@ class AAttn(nn.Module):
         self.pe = Conv(all_head_dim, dim, 7, 1, 3, g=dim, act=False)
 
     def forward(self, x):
-        """Processes the input tensor 'x' through the area-attention."""
+        """
+        Process the input tensor through the area-attention.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor after area-attention.
+        """
         B, C, H, W = x.shape
         N = H * W
 
@@ -1260,11 +1746,7 @@ class ABlock(nn.Module):
 
     def __init__(self, dim, num_heads, mlp_ratio=1.2, area=1):
         """
-        Initializes an Area-attention block module for efficient feature extraction in YOLO models.
-
-        This module implements an area-attention mechanism combined with a feed-forward network for processing feature
-        maps. It uses a novel area-based attention approach that is more efficient than traditional self-attention
-        while maintaining effectiveness.
+        Initialize an Area-attention block module.
 
         Args:
             dim (int): Number of input channels.
@@ -1281,14 +1763,27 @@ class ABlock(nn.Module):
         self.apply(self._init_weights)
 
     def _init_weights(self, m):
-        """Initialize weights using a truncated normal distribution."""
+        """
+        Initialize weights using a truncated normal distribution.
+
+        Args:
+            m (nn.Module): Module to initialize.
+        """
         if isinstance(m, nn.Conv2d):
             nn.init.trunc_normal_(m.weight, std=0.02)
             if m.bias is not None:
                 nn.init.constant_(m.bias, 0)
 
     def forward(self, x):
-        """Forward pass through ABlock, applying area-attention and feed-forward layers to the input tensor."""
+        """
+        Forward pass through ABlock.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor after area-attention and feed-forward processing.
+        """
         x = x + self.attn(x)
         return x + self.mlp(x)
 
@@ -1319,7 +1814,7 @@ class A2C2f(nn.Module):
 
     def __init__(self, c1, c2, n=1, a2=True, area=1, residual=False, mlp_ratio=2.0, e=0.5, g=1, shortcut=True):
         """
-        Area-Attention C2f module for enhanced feature extraction with area-based attention mechanisms.
+        Initialize Area-Attention C2f module.
 
         Args:
             c1 (int): Number of input channels.
@@ -1349,7 +1844,15 @@ class A2C2f(nn.Module):
         )
 
     def forward(self, x):
-        """Forward pass through R-ELAN layer."""
+        """
+        Forward pass through A2C2f layer.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            (torch.Tensor): Output tensor after processing.
+        """
         y = [self.cv1(x)]
         y.extend(m(y[-1]) for m in self.m)
         y = self.cv2(torch.cat(y, 1))