Reference for `ultralytics/nn/modules/block.py`

Note

This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/nn/modules/block.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!

ultralytics.nn.modules.block.DFL

DFL(c1=16)

Bases: Module

Integral module of Distribution Focal Loss (DFL).

Proposed in Generalized Focal Loss https://ieeexplore.ieee.org/document/9792391

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1=16):
    """Initialize a convolutional layer with a given number of input channels."""
    super().__init__()
    self.conv = nn.Conv2d(c1, 1, 1, bias=False).requires_grad_(False)
    x = torch.arange(c1, dtype=torch.float)
    self.conv.weight.data[:] = nn.Parameter(x.view(1, c1, 1, 1))
    self.c1 = c1

forward

forward(x)

Apply the DFL module to input tensor and return transformed output.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """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)

ultralytics.nn.modules.block.Proto

Proto(c1, c_=256, c2=32)

Bases: Module

YOLOv8 mask Proto module for segmentation models.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c_`	`int`	Intermediate channels.	`256`
`c2`	`int`	Output channels (number of protos).	`32`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c_=256, c2=32):
    """
    Initialize the YOLOv8 mask Proto module with specified number of protos and 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)
    self.upsample = nn.ConvTranspose2d(c_, c_, 2, 2, 0, bias=True)  # nn.Upsample(scale_factor=2, mode='nearest')
    self.cv2 = Conv(c_, c_, k=3)
    self.cv3 = Conv(c_, c2)

forward

forward(x)

Perform a forward pass through layers using an upsampled input image.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Perform a forward pass through layers using an upsampled input image."""
    return self.cv3(self.cv2(self.upsample(self.cv1(x))))

ultralytics.nn.modules.block.HGStem

HGStem(c1, cm, c2)

Bases: Module

StemBlock of PPHGNetV2 with 5 convolutions and one maxpool2d.

https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/backbones/hgnet_v2.py

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`cm`	`int`	Middle channels.	required
`c2`	`int`	Output channels.	required

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, cm, c2):
    """
    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())
    self.stem2b = Conv(cm // 2, cm, 2, 1, 0, act=nn.ReLU())
    self.stem3 = Conv(cm * 2, cm, 3, 2, act=nn.ReLU())
    self.stem4 = Conv(cm, c2, 1, 1, act=nn.ReLU())
    self.pool = nn.MaxPool2d(kernel_size=2, stride=1, padding=0, ceil_mode=True)

forward

forward(x)

Forward pass of a PPHGNetV2 backbone layer.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Forward pass of a PPHGNetV2 backbone layer."""
    x = self.stem1(x)
    x = F.pad(x, [0, 1, 0, 1])
    x2 = self.stem2a(x)
    x2 = F.pad(x2, [0, 1, 0, 1])
    x2 = self.stem2b(x2)
    x1 = self.pool(x)
    x = torch.cat([x1, x2], dim=1)
    x = self.stem3(x)
    x = self.stem4(x)
    return x

ultralytics.nn.modules.block.HGBlock

HGBlock(c1, cm, c2, k=3, n=6, lightconv=False, shortcut=False, act=nn.ReLU())

Bases: Module

HG_Block of PPHGNetV2 with 2 convolutions and LightConv.

https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/backbones/hgnet_v2.py

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`cm`	`int`	Middle channels.	required
`c2`	`int`	Output channels.	required
`k`	`int`	Kernel size.	`3`
`n`	`int`	Number of LightConv or Conv blocks.	`6`
`lightconv`	`bool`	Whether to use LightConv.	`False`
`shortcut`	`bool`	Whether to use shortcut connection.	`False`
`act`	`Module`	Activation function.	`ReLU()`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, cm, c2, k=3, n=6, lightconv=False, shortcut=False, act=nn.ReLU()):
    """
    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))
    self.sc = Conv(c1 + n * cm, c2 // 2, 1, 1, act=act)  # squeeze conv
    self.ec = Conv(c2 // 2, c2, 1, 1, act=act)  # excitation conv
    self.add = shortcut and c1 == c2

forward

forward(x)

Forward pass of a PPHGNetV2 backbone layer.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Forward pass of a PPHGNetV2 backbone layer."""
    y = [x]
    y.extend(m(y[-1]) for m in self.m)
    y = self.ec(self.sc(torch.cat(y, 1)))
    return y + x if self.add else y

ultralytics.nn.modules.block.SPP

SPP(c1, c2, k=(5, 9, 13))

Bases: Module

Spatial Pyramid Pooling (SPP) layer https://arxiv.org/abs/1406.4729.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`k`	`Tuple[int, int, int]`	Kernel sizes for max pooling.	`(5, 9, 13)`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, k=(5, 9, 13)):
    """
    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)
    self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
    self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])

forward

forward(x)

Forward pass of the SPP layer, performing spatial pyramid pooling.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Forward pass of the SPP layer, performing spatial pyramid pooling."""
    x = self.cv1(x)
    return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))

ultralytics.nn.modules.block.SPPF

SPPF(c1, c2, k=5)

Bases: Module

Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`k`	`int`	Kernel size.	`5`

Notes

This module is equivalent to SPP(k=(5, 9, 13)).

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, k=5):
    """
    Initialize the SPPF layer with given input/output channels and kernel size.

    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
    self.cv1 = Conv(c1, c_, 1, 1)
    self.cv2 = Conv(c_ * 4, c2, 1, 1)
    self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)

forward

forward(x)

Apply sequential pooling operations to input and return concatenated feature maps.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """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))

ultralytics.nn.modules.block.C1

C1(c1, c2, n=1)

Bases: Module

CSP Bottleneck with 1 convolution.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`n`	`int`	Number of convolutions.	`1`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=1):
    """
    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)))

forward

forward(x)

Apply convolution and residual connection to input tensor.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Apply convolution and residual connection to input tensor."""
    y = self.cv1(x)
    return self.m(y) + y

ultralytics.nn.modules.block.C2

C2(c1, c2, n=1, shortcut=True, g=1, e=0.5)

Bases: Module

CSP Bottleneck with 2 convolutions.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`n`	`int`	Number of Bottleneck blocks.	`1`
`shortcut`	`bool`	Whether to use shortcut connections.	`True`
`g`	`int`	Groups for convolutions.	`1`
`e`	`float`	Expansion ratio.	`0.5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
    """
    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)
    self.cv2 = Conv(2 * self.c, c2, 1)  # optional act=FReLU(c2)
    # self.attention = ChannelAttention(2 * self.c)  # or SpatialAttention()
    self.m = nn.Sequential(*(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n)))

forward

forward(x)

Forward pass through the CSP bottleneck with 2 convolutions.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Forward pass through the CSP bottleneck with 2 convolutions."""
    a, b = self.cv1(x).chunk(2, 1)
    return self.cv2(torch.cat((self.m(a), b), 1))

ultralytics.nn.modules.block.C2f

C2f(c1, c2, n=1, shortcut=False, g=1, e=0.5)

Bases: Module

Faster Implementation of CSP Bottleneck with 2 convolutions.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`n`	`int`	Number of Bottleneck blocks.	`1`
`shortcut`	`bool`	Whether to use shortcut connections.	`False`
`g`	`int`	Groups for convolutions.	`1`
`e`	`float`	Expansion ratio.	`0.5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
    """
    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)
    self.cv2 = Conv((2 + n) * self.c, c2, 1)  # optional act=FReLU(c2)
    self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))

forward

forward(x)

Forward pass through C2f layer.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Forward pass through C2f layer."""
    y = list(self.cv1(x).chunk(2, 1))
    y.extend(m(y[-1]) for m in self.m)
    return self.cv2(torch.cat(y, 1))

forward_split

forward_split(x)

Forward pass using split() instead of chunk().

Source code in ultralytics/nn/modules/block.py

def forward_split(self, x):
    """Forward pass using split() instead of chunk()."""
    y = self.cv1(x).split((self.c, self.c), 1)
    y = [y[0], y[1]]
    y.extend(m(y[-1]) for m in self.m)
    return self.cv2(torch.cat(y, 1))

ultralytics.nn.modules.block.C3

C3(c1, c2, n=1, shortcut=True, g=1, e=0.5)

Bases: Module

CSP Bottleneck with 3 convolutions.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`n`	`int`	Number of Bottleneck blocks.	`1`
`shortcut`	`bool`	Whether to use shortcut connections.	`True`
`g`	`int`	Groups for convolutions.	`1`
`e`	`float`	Expansion ratio.	`0.5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
    """
    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)
    self.cv2 = Conv(c1, c_, 1, 1)
    self.cv3 = Conv(2 * c_, c2, 1)  # optional act=FReLU(c2)
    self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, k=((1, 1), (3, 3)), e=1.0) for _ in range(n)))

forward

forward(x)

Forward pass through the CSP bottleneck with 3 convolutions.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Forward pass through the CSP bottleneck with 3 convolutions."""
    return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))

ultralytics.nn.modules.block.C3x

C3x(c1, c2, n=1, shortcut=True, g=1, e=0.5)

Bases: C3

C3 module with cross-convolutions.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`n`	`int`	Number of Bottleneck blocks.	`1`
`shortcut`	`bool`	Whether to use shortcut connections.	`True`
`g`	`int`	Groups for convolutions.	`1`
`e`	`float`	Expansion ratio.	`0.5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
    """
    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)))

ultralytics.nn.modules.block.RepC3

RepC3(c1, c2, n=3, e=1.0)

Bases: Module

Rep C3.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`n`	`int`	Number of RepConv blocks.	`3`
`e`	`float`	Expansion ratio.	`1.0`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=3, e=1.0):
    """
    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)
    self.cv2 = Conv(c1, c_, 1, 1)
    self.m = nn.Sequential(*[RepConv(c_, c_) for _ in range(n)])
    self.cv3 = Conv(c_, c2, 1, 1) if c_ != c2 else nn.Identity()

forward

forward(x)

Forward pass of RepC3 module.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Forward pass of RepC3 module."""
    return self.cv3(self.m(self.cv1(x)) + self.cv2(x))

ultralytics.nn.modules.block.C3TR

C3TR(c1, c2, n=1, shortcut=True, g=1, e=0.5)

Bases: C3

C3 module with TransformerBlock().

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`n`	`int`	Number of Transformer blocks.	`1`
`shortcut`	`bool`	Whether to use shortcut connections.	`True`
`g`	`int`	Groups for convolutions.	`1`
`e`	`float`	Expansion ratio.	`0.5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
    """
    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)

ultralytics.nn.modules.block.C3Ghost

C3Ghost(c1, c2, n=1, shortcut=True, g=1, e=0.5)

Bases: C3

C3 module with GhostBottleneck().

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`n`	`int`	Number of Ghost bottleneck blocks.	`1`
`shortcut`	`bool`	Whether to use shortcut connections.	`True`
`g`	`int`	Groups for convolutions.	`1`
`e`	`float`	Expansion ratio.	`0.5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
    """
    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)))

ultralytics.nn.modules.block.GhostBottleneck

GhostBottleneck(c1, c2, k=3, s=1)

Bases: Module

Ghost Bottleneck https://github.com/huawei-noah/ghostnet.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`k`	`int`	Kernel size.	`3`
`s`	`int`	Stride.	`1`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, k=3, s=1):
    """
    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(
        GhostConv(c1, c_, 1, 1),  # pw
        DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(),  # dw
        GhostConv(c_, c2, 1, 1, act=False),  # pw-linear
    )
    self.shortcut = (
        nn.Sequential(DWConv(c1, c1, k, s, act=False), Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity()
    )

forward

forward(x)

Apply skip connection and concatenation to input tensor.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Apply skip connection and concatenation to input tensor."""
    return self.conv(x) + self.shortcut(x)

ultralytics.nn.modules.block.Bottleneck

Bottleneck(c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5)

Bases: Module

Standard bottleneck.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`shortcut`	`bool`	Whether to use shortcut connection.	`True`
`g`	`int`	Groups for convolutions.	`1`
`k`	`Tuple[int, int]`	Kernel sizes for convolutions.	`(3, 3)`
`e`	`float`	Expansion ratio.	`0.5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):
    """
    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)
    self.cv2 = Conv(c_, c2, k[1], 1, g=g)
    self.add = shortcut and c1 == c2

forward

forward(x)

Apply bottleneck with optional shortcut connection.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Apply bottleneck with optional shortcut connection."""
    return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))

ultralytics.nn.modules.block.BottleneckCSP

BottleneckCSP(c1, c2, n=1, shortcut=True, g=1, e=0.5)

Bases: Module

CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`n`	`int`	Number of Bottleneck blocks.	`1`
`shortcut`	`bool`	Whether to use shortcut connections.	`True`
`g`	`int`	Groups for convolutions.	`1`
`e`	`float`	Expansion ratio.	`0.5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
    """
    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)
    self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
    self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
    self.cv4 = Conv(2 * c_, c2, 1, 1)
    self.bn = nn.BatchNorm2d(2 * c_)  # applied to cat(cv2, cv3)
    self.act = nn.SiLU()
    self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))

forward

forward(x)

Apply CSP bottleneck with 3 convolutions.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """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))))

ultralytics.nn.modules.block.ResNetBlock

ResNetBlock(c1, c2, s=1, e=4)

Bases: Module

ResNet block with standard convolution layers.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`s`	`int`	Stride.	`1`
`e`	`int`	Expansion ratio.	`4`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, s=1, e=4):
    """
    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)
    self.cv2 = Conv(c2, c2, k=3, s=s, p=1, act=True)
    self.cv3 = Conv(c2, c3, k=1, act=False)
    self.shortcut = nn.Sequential(Conv(c1, c3, k=1, s=s, act=False)) if s != 1 or c1 != c3 else nn.Identity()

forward

forward(x)

Forward pass through the ResNet block.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Forward pass through the ResNet block."""
    return F.relu(self.cv3(self.cv2(self.cv1(x))) + self.shortcut(x))

ultralytics.nn.modules.block.ResNetLayer

ResNetLayer(c1, c2, s=1, is_first=False, n=1, e=4)

Bases: Module

ResNet layer with multiple ResNet blocks.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`s`	`int`	Stride.	`1`
`is_first`	`bool`	Whether this is the first layer.	`False`
`n`	`int`	Number of ResNet blocks.	`1`
`e`	`int`	Expansion ratio.	`4`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, s=1, is_first=False, n=1, e=4):
    """
    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

    if self.is_first:
        self.layer = nn.Sequential(
            Conv(c1, c2, k=7, s=2, p=3, act=True), nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )
    else:
        blocks = [ResNetBlock(c1, c2, s, e=e)]
        blocks.extend([ResNetBlock(e * c2, c2, 1, e=e) for _ in range(n - 1)])
        self.layer = nn.Sequential(*blocks)

forward

forward(x)

Forward pass through the ResNet layer.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Forward pass through the ResNet layer."""
    return self.layer(x)

ultralytics.nn.modules.block.MaxSigmoidAttnBlock

MaxSigmoidAttnBlock(c1, c2, nh=1, ec=128, gc=512, scale=False)

Bases: Module

Max Sigmoid attention block.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`nh`	`int`	Number of heads.	`1`
`ec`	`int`	Embedding channels.	`128`
`gc`	`int`	Guide channels.	`512`
`scale`	`bool`	Whether to use learnable scale parameter.	`False`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, nh=1, ec=128, gc=512, scale=False):
    """
    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
    self.ec = Conv(c1, ec, k=1, act=False) if c1 != ec else None
    self.gl = nn.Linear(gc, ec)
    self.bias = nn.Parameter(torch.zeros(nh))
    self.proj_conv = Conv(c1, c2, k=3, s=1, act=False)
    self.scale = nn.Parameter(torch.ones(1, nh, 1, 1)) if scale else 1.0

forward

forward(x, guide)

Forward pass of MaxSigmoidAttnBlock.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor.	required
`guide`	`Tensor`	Guide tensor.	required

Returns:

Type	Description
`Tensor`	Output tensor after attention.

Source code in ultralytics/nn/modules/block.py

def forward(self, x, guide):
    """
    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)
    guide = guide.view(bs, -1, self.nh, self.hc)
    embed = self.ec(x) if self.ec is not None else x
    embed = embed.view(bs, self.nh, self.hc, h, w)

    aw = torch.einsum("bmchw,bnmc->bmhwn", embed, guide)
    aw = aw.max(dim=-1)[0]
    aw = aw / (self.hc**0.5)
    aw = aw + self.bias[None, :, None, None]
    aw = aw.sigmoid() * self.scale

    x = self.proj_conv(x)
    x = x.view(bs, self.nh, -1, h, w)
    x = x * aw.unsqueeze(2)
    return x.view(bs, -1, h, w)

ultralytics.nn.modules.block.C2fAttn

C2fAttn(c1, c2, n=1, ec=128, nh=1, gc=512, shortcut=False, g=1, e=0.5)

Bases: Module

C2f module with an additional attn module.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`n`	`int`	Number of Bottleneck blocks.	`1`
`ec`	`int`	Embedding channels for attention.	`128`
`nh`	`int`	Number of heads for attention.	`1`
`gc`	`int`	Guide channels for attention.	`512`
`shortcut`	`bool`	Whether to use shortcut connections.	`False`
`g`	`int`	Groups for convolutions.	`1`
`e`	`float`	Expansion ratio.	`0.5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=1, ec=128, nh=1, gc=512, shortcut=False, g=1, e=0.5):
    """
    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)
    self.cv2 = Conv((3 + n) * self.c, c2, 1)  # optional act=FReLU(c2)
    self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))
    self.attn = MaxSigmoidAttnBlock(self.c, self.c, gc=gc, ec=ec, nh=nh)

forward

forward(x, guide)

Forward pass through C2f layer with attention.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor.	required
`guide`	`Tensor`	Guide tensor for attention.	required

Returns:

Type	Description
`Tensor`	Output tensor after processing.

Source code in ultralytics/nn/modules/block.py

def forward(self, x, guide):
    """
    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))

forward_split

forward_split(x, guide)

Forward pass using split() instead of chunk().

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor.	required
`guide`	`Tensor`	Guide tensor for attention.	required

Returns:

Type	Description
`Tensor`	Output tensor after processing.

Source code in ultralytics/nn/modules/block.py

def forward_split(self, x, guide):
    """
    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))
    return self.cv2(torch.cat(y, 1))

ultralytics.nn.modules.block.ImagePoolingAttn

ImagePoolingAttn(ec=256, ch=(), ct=512, nh=8, k=3, scale=False)

Bases: Module

ImagePoolingAttn: Enhance the text embeddings with image-aware information.

Parameters:

Name	Type	Description	Default
`ec`	`int`	Embedding channels.	`256`
`ch`	`tuple`	Channel dimensions for feature maps.	`()`
`ct`	`int`	Channel dimension for text embeddings.	`512`
`nh`	`int`	Number of attention heads.	`8`
`k`	`int`	Kernel size for pooling.	`3`
`scale`	`bool`	Whether to use learnable scale parameter.	`False`

Source code in ultralytics/nn/modules/block.py

def __init__(self, ec=256, ch=(), ct=512, nh=8, k=3, scale=False):
    """
    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)
    self.query = nn.Sequential(nn.LayerNorm(ct), nn.Linear(ct, ec))
    self.key = nn.Sequential(nn.LayerNorm(ec), nn.Linear(ec, ec))
    self.value = nn.Sequential(nn.LayerNorm(ec), nn.Linear(ec, ec))
    self.proj = nn.Linear(ec, ct)
    self.scale = nn.Parameter(torch.tensor([0.0]), requires_grad=True) if scale else 1.0
    self.projections = nn.ModuleList([nn.Conv2d(in_channels, ec, kernel_size=1) for in_channels in ch])
    self.im_pools = nn.ModuleList([nn.AdaptiveMaxPool2d((k, k)) for _ in range(nf)])
    self.ec = ec
    self.nh = nh
    self.nf = nf
    self.hc = ec // nh
    self.k = k

forward

forward(x, text)

Forward pass of ImagePoolingAttn.

Parameters:

Name	Type	Description	Default
`x`	`List[Tensor]`	List of input feature maps.	required
`text`	`Tensor`	Text embeddings.	required

Returns:

Type	Description
`Tensor`	Enhanced text embeddings.

Source code in ultralytics/nn/modules/block.py

def forward(self, x, text):
    """
    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
    x = [pool(proj(x)).view(bs, -1, num_patches) for (x, proj, pool) in zip(x, self.projections, self.im_pools)]
    x = torch.cat(x, dim=-1).transpose(1, 2)
    q = self.query(text)
    k = self.key(x)
    v = self.value(x)

    # q = q.reshape(1, text.shape[1], self.nh, self.hc).repeat(bs, 1, 1, 1)
    q = q.reshape(bs, -1, self.nh, self.hc)
    k = k.reshape(bs, -1, self.nh, self.hc)
    v = v.reshape(bs, -1, self.nh, self.hc)

    aw = torch.einsum("bnmc,bkmc->bmnk", q, k)
    aw = aw / (self.hc**0.5)
    aw = F.softmax(aw, dim=-1)

    x = torch.einsum("bmnk,bkmc->bnmc", aw, v)
    x = self.proj(x.reshape(bs, -1, self.ec))
    return x * self.scale + text

ultralytics.nn.modules.block.ContrastiveHead

ContrastiveHead()

Bases: Module

Implements contrastive learning head for region-text similarity in vision-language models.

Source code in ultralytics/nn/modules/block.py

def __init__(self):
    """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())

forward

forward(x, w)

Forward function of contrastive learning.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Image features.	required
`w`	`Tensor`	Text features.	required

Returns:

Type	Description
`Tensor`	Similarity scores.

Source code in ultralytics/nn/modules/block.py

def forward(self, x, w):
    """
    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)
    return x * self.logit_scale.exp() + self.bias

ultralytics.nn.modules.block.BNContrastiveHead

BNContrastiveHead(embed_dims: int)

Bases: Module

Batch Norm Contrastive Head for YOLO-World using batch norm instead of l2-normalization.

Parameters:

Name	Type	Description	Default
`embed_dims`	`int`	Embed dimensions of text and image features.	required

Parameters:

Name	Type	Description	Default
`embed_dims`	`int`	Embedding dimensions for features.	required

Source code in ultralytics/nn/modules/block.py

def __init__(self, embed_dims: int):
    """
    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
    self.bias = nn.Parameter(torch.tensor([-10.0]))
    # use -1.0 is more stable
    self.logit_scale = nn.Parameter(-1.0 * torch.ones([]))

forward

forward(x, w)

Forward function of contrastive learning with batch normalization.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Image features.	required
`w`	`Tensor`	Text features.	required

Returns:

Type	Description
`Tensor`	Similarity scores.

Source code in ultralytics/nn/modules/block.py

def forward(self, x, w):
    """
    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)
    return x * self.logit_scale.exp() + self.bias

ultralytics.nn.modules.block.RepBottleneck

RepBottleneck(c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5)

Bases: Bottleneck

Rep bottleneck.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`shortcut`	`bool`	Whether to use shortcut connection.	`True`
`g`	`int`	Groups for convolutions.	`1`
`k`	`Tuple[int, int]`	Kernel sizes for convolutions.	`(3, 3)`
`e`	`float`	Expansion ratio.	`0.5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):
    """
    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)

ultralytics.nn.modules.block.RepCSP

RepCSP(c1, c2, n=1, shortcut=True, g=1, e=0.5)

Bases: C3

Repeatable Cross Stage Partial Network (RepCSP) module for efficient feature extraction.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`n`	`int`	Number of RepBottleneck blocks.	`1`
`shortcut`	`bool`	Whether to use shortcut connections.	`True`
`g`	`int`	Groups for convolutions.	`1`
`e`	`float`	Expansion ratio.	`0.5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
    """
    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)))

ultralytics.nn.modules.block.RepNCSPELAN4

RepNCSPELAN4(c1, c2, c3, c4, n=1)

Bases: Module

CSP-ELAN.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`c3`	`int`	Intermediate channels.	required
`c4`	`int`	Intermediate channels for RepCSP.	required
`n`	`int`	Number of RepCSP blocks.	`1`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, c3, c4, n=1):
    """
    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)
    self.cv2 = nn.Sequential(RepCSP(c3 // 2, c4, n), Conv(c4, c4, 3, 1))
    self.cv3 = nn.Sequential(RepCSP(c4, c4, n), Conv(c4, c4, 3, 1))
    self.cv4 = Conv(c3 + (2 * c4), c2, 1, 1)

forward

forward(x)

Forward pass through RepNCSPELAN4 layer.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Forward pass through RepNCSPELAN4 layer."""
    y = list(self.cv1(x).chunk(2, 1))
    y.extend((m(y[-1])) for m in [self.cv2, self.cv3])
    return self.cv4(torch.cat(y, 1))

forward_split

forward_split(x)

Forward pass using split() instead of chunk().

Source code in ultralytics/nn/modules/block.py

def forward_split(self, x):
    """Forward pass using split() instead of chunk()."""
    y = list(self.cv1(x).split((self.c, self.c), 1))
    y.extend(m(y[-1]) for m in [self.cv2, self.cv3])
    return self.cv4(torch.cat(y, 1))

ultralytics.nn.modules.block.ELAN1

ELAN1(c1, c2, c3, c4)

Bases: RepNCSPELAN4

ELAN1 module with 4 convolutions.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`c3`	`int`	Intermediate channels.	required
`c4`	`int`	Intermediate channels for convolutions.	required

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, c3, c4):
    """
    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)
    self.cv2 = Conv(c3 // 2, c4, 3, 1)
    self.cv3 = Conv(c4, c4, 3, 1)
    self.cv4 = Conv(c3 + (2 * c4), c2, 1, 1)

ultralytics.nn.modules.block.AConv

AConv(c1, c2)

Bases: Module

AConv.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2):
    """
    Initialize AConv module.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
    """
    super().__init__()
    self.cv1 = Conv(c1, c2, 3, 2, 1)

forward

forward(x)

Forward pass through AConv layer.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Forward pass through AConv layer."""
    x = torch.nn.functional.avg_pool2d(x, 2, 1, 0, False, True)
    return self.cv1(x)

ultralytics.nn.modules.block.ADown

ADown(c1, c2)

Bases: Module

ADown.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, 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)
    self.cv2 = Conv(c1 // 2, self.c, 1, 1, 0)

forward

forward(x)

Forward pass through ADown layer.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Forward pass through ADown layer."""
    x = torch.nn.functional.avg_pool2d(x, 2, 1, 0, False, True)
    x1, x2 = x.chunk(2, 1)
    x1 = self.cv1(x1)
    x2 = torch.nn.functional.max_pool2d(x2, 3, 2, 1)
    x2 = self.cv2(x2)
    return torch.cat((x1, x2), 1)

ultralytics.nn.modules.block.SPPELAN

SPPELAN(c1, c2, c3, k=5)

Bases: Module

SPP-ELAN.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`c3`	`int`	Intermediate channels.	required
`k`	`int`	Kernel size for max pooling.	`5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, c3, k=5):
    """
    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)
    self.cv2 = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
    self.cv3 = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
    self.cv4 = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
    self.cv5 = Conv(4 * c3, c2, 1, 1)

forward

forward(x)

Forward pass through SPPELAN layer.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Forward pass through SPPELAN layer."""
    y = [self.cv1(x)]
    y.extend(m(y[-1]) for m in [self.cv2, self.cv3, self.cv4])
    return self.cv5(torch.cat(y, 1))

ultralytics.nn.modules.block.CBLinear

CBLinear(c1, c2s, k=1, s=1, p=None, g=1)

Bases: Module

CBLinear.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2s`	`List[int]`	List of output channel sizes.	required
`k`	`int`	Kernel size.	`1`
`s`	`int`	Stride.	`1`
`p`	`int \| None`	Padding.	`None`
`g`	`int`	Groups.	`1`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2s, k=1, s=1, p=None, g=1):
    """
    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)

forward

forward(x)

Forward pass through CBLinear layer.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """Forward pass through CBLinear layer."""
    return self.conv(x).split(self.c2s, dim=1)

ultralytics.nn.modules.block.CBFuse

CBFuse(idx)

Bases: Module

CBFuse.

Parameters:

Name	Type	Description	Default
`idx`	`List[int]`	Indices for feature selection.	required

Source code in ultralytics/nn/modules/block.py

def __init__(self, idx):
    """
    Initialize CBFuse module.

    Args:
        idx (List[int]): Indices for feature selection.
    """
    super().__init__()
    self.idx = idx

forward

forward(xs)

Forward pass through CBFuse layer.

Parameters:

Name	Type	Description	Default
`xs`	`List[Tensor]`	List of input tensors.	required

Returns:

Type	Description
`Tensor`	Fused output tensor.

Source code in ultralytics/nn/modules/block.py

def forward(self, xs):
    """
    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)

ultralytics.nn.modules.block.C3f

C3f(c1, c2, n=1, shortcut=False, g=1, e=0.5)

Bases: Module

Faster Implementation of CSP Bottleneck with 2 convolutions.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`n`	`int`	Number of Bottleneck blocks.	`1`
`shortcut`	`bool`	Whether to use shortcut connections.	`False`
`g`	`int`	Groups for convolutions.	`1`
`e`	`float`	Expansion ratio.	`0.5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
    """
    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
    self.cv1 = Conv(c1, c_, 1, 1)
    self.cv2 = Conv(c1, c_, 1, 1)
    self.cv3 = Conv((2 + n) * c_, c2, 1)  # optional act=FReLU(c2)
    self.m = nn.ModuleList(Bottleneck(c_, c_, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))

forward

forward(x)

Forward pass through C3f layer.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """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))

ultralytics.nn.modules.block.C3k2

C3k2(c1, c2, n=1, c3k=False, e=0.5, g=1, shortcut=True)

Bases: C2f

Faster Implementation of CSP Bottleneck with 2 convolutions.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`n`	`int`	Number of blocks.	`1`
`c3k`	`bool`	Whether to use C3k blocks.	`False`
`e`	`float`	Expansion ratio.	`0.5`
`g`	`int`	Groups for convolutions.	`1`
`shortcut`	`bool`	Whether to use shortcut connections.	`True`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=1, c3k=False, e=0.5, g=1, shortcut=True):
    """
    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)
    )

ultralytics.nn.modules.block.C3k

C3k(c1, c2, n=1, shortcut=True, g=1, e=0.5, k=3)

Bases: C3

C3k is a CSP bottleneck module with customizable kernel sizes for feature extraction in neural networks.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`n`	`int`	Number of Bottleneck blocks.	`1`
`shortcut`	`bool`	Whether to use shortcut connections.	`True`
`g`	`int`	Groups for convolutions.	`1`
`e`	`float`	Expansion ratio.	`0.5`
`k`	`int`	Kernel size.	`3`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, k=3):
    """
    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)))
    self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, k=(k, k), e=1.0) for _ in range(n)))

ultralytics.nn.modules.block.RepVGGDW

RepVGGDW(ed)

Bases: Module

RepVGGDW is a class that represents a depth wise separable convolutional block in RepVGG architecture.

Parameters:

Name	Type	Description	Default
`ed`	`int`	Input and output channels.	required

Source code in ultralytics/nn/modules/block.py

def __init__(self, ed) -> None:
    """
    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)
    self.dim = ed
    self.act = nn.SiLU()

forward

forward(x)

Perform a forward pass of the RepVGGDW block.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor.	required

Returns:

Type	Description
`Tensor`	Output tensor after applying the depth wise separable convolution.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """
    Perform a forward pass of the RepVGGDW block.

    Args:
        x (torch.Tensor): Input tensor.

    Returns:
        (torch.Tensor): Output tensor after applying the depth wise separable convolution.
    """
    return self.act(self.conv(x) + self.conv1(x))

forward_fuse

forward_fuse(x)

Perform a forward pass of the RepVGGDW block without fusing the convolutions.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor.	required

Returns:

Type	Description
`Tensor`	Output tensor after applying the depth wise separable convolution.

Source code in ultralytics/nn/modules/block.py

def forward_fuse(self, x):
    """
    Perform a forward pass of the RepVGGDW block without fusing the convolutions.

    Args:
        x (torch.Tensor): Input tensor.

    Returns:
        (torch.Tensor): Output tensor after applying the depth wise separable convolution.
    """
    return self.act(self.conv(x))

fuse

fuse()

Fuse the convolutional layers in the RepVGGDW block.

This method fuses the convolutional layers and updates the weights and biases accordingly.

Source code in ultralytics/nn/modules/block.py

@torch.no_grad()
def fuse(self):
    """
    Fuse the convolutional layers in the RepVGGDW block.

    This method fuses the convolutional layers and updates the weights and biases accordingly.
    """
    conv = fuse_conv_and_bn(self.conv.conv, self.conv.bn)
    conv1 = fuse_conv_and_bn(self.conv1.conv, self.conv1.bn)

    conv_w = conv.weight
    conv_b = conv.bias
    conv1_w = conv1.weight
    conv1_b = conv1.bias

    conv1_w = torch.nn.functional.pad(conv1_w, [2, 2, 2, 2])

    final_conv_w = conv_w + conv1_w
    final_conv_b = conv_b + conv1_b

    conv.weight.data.copy_(final_conv_w)
    conv.bias.data.copy_(final_conv_b)

    self.conv = conv
    del self.conv1

ultralytics.nn.modules.block.CIB

CIB(c1, c2, shortcut=True, e=0.5, lk=False)

Bases: Module

Conditional Identity Block (CIB) module.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Number of input channels.	required
`c2`	`int`	Number of output channels.	required
`shortcut`	`bool`	Whether to add a shortcut connection. Defaults to True.	`True`
`e`	`float`	Scaling factor for the hidden channels. Defaults to 0.5.	`0.5`
`lk`	`bool`	Whether to use RepVGGDW for the third convolutional layer. Defaults to False.	`False`

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`shortcut`	`bool`	Whether to use shortcut connection.	`True`
`e`	`float`	Expansion ratio.	`0.5`
`lk`	`bool`	Whether to use RepVGGDW.	`False`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, shortcut=True, e=0.5, lk=False):
    """
    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(
        Conv(c1, c1, 3, g=c1),
        Conv(c1, 2 * c_, 1),
        RepVGGDW(2 * c_) if lk else Conv(2 * c_, 2 * c_, 3, g=2 * c_),
        Conv(2 * c_, c2, 1),
        Conv(c2, c2, 3, g=c2),
    )

    self.add = shortcut and c1 == c2

forward

forward(x)

Forward pass of the CIB module.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor.	required

Returns:

Type	Description
`Tensor`	Output tensor.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """
    Forward pass of the CIB module.

    Args:
        x (torch.Tensor): Input tensor.

    Returns:
        (torch.Tensor): Output tensor.
    """
    return x + self.cv1(x) if self.add else self.cv1(x)

ultralytics.nn.modules.block.C2fCIB

C2fCIB(c1, c2, n=1, shortcut=False, lk=False, g=1, e=0.5)

Bases: C2f

C2fCIB class represents a convolutional block with C2f and CIB modules.

Parameters:

Name	Type	Description	Default
`c1`	`int`	Number of input channels.	required
`c2`	`int`	Number of output channels.	required
`n`	`int`	Number of CIB modules to stack. Defaults to 1.	`1`
`shortcut`	`bool`	Whether to use shortcut connection. Defaults to False.	`False`
`lk`	`bool`	Whether to use local key connection. Defaults to False.	`False`
`g`	`int`	Number of groups for grouped convolution. Defaults to 1.	`1`
`e`	`float`	Expansion ratio for CIB modules. Defaults to 0.5.	`0.5`

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`n`	`int`	Number of CIB modules.	`1`
`shortcut`	`bool`	Whether to use shortcut connection.	`False`
`lk`	`bool`	Whether to use local key connection.	`False`
`g`	`int`	Groups for convolutions.	`1`
`e`	`float`	Expansion ratio.	`0.5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=1, shortcut=False, lk=False, g=1, e=0.5):
    """
    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))

ultralytics.nn.modules.block.Attention

Attention(dim, num_heads=8, attn_ratio=0.5)

Bases: Module

Attention module that performs self-attention on the input tensor.

Parameters:

Name	Type	Description	Default
`dim`	`int`	The input tensor dimension.	required
`num_heads`	`int`	The number of attention heads.	`8`
`attn_ratio`	`float`	The ratio of the attention key dimension to the head dimension.	`0.5`

Attributes:

Name	Type	Description
`num_heads`	`int`	The number of attention heads.
`head_dim`	`int`	The dimension of each attention head.
`key_dim`	`int`	The dimension of the attention key.
`scale`	`float`	The scaling factor for the attention scores.
`qkv`	`Conv`	Convolutional layer for computing the query, key, and value.
`proj`	`Conv`	Convolutional layer for projecting the attended values.
`pe`	`Conv`	Convolutional layer for positional encoding.

Parameters:

Name	Type	Description	Default
`dim`	`int`	Input dimension.	required
`num_heads`	`int`	Number of attention heads.	`8`
`attn_ratio`	`float`	Attention ratio for key dimension.	`0.5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, dim, num_heads=8, attn_ratio=0.5):
    """
    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
    self.key_dim = int(self.head_dim * attn_ratio)
    self.scale = self.key_dim**-0.5
    nh_kd = self.key_dim * num_heads
    h = dim + nh_kd * 2
    self.qkv = Conv(dim, h, 1, act=False)
    self.proj = Conv(dim, dim, 1, act=False)
    self.pe = Conv(dim, dim, 3, 1, g=dim, act=False)

forward

forward(x)

Forward pass of the Attention module.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	The input tensor.	required

Returns:

Type	Description
`Tensor`	The output tensor after self-attention.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """
    Forward pass of the Attention module.

    Args:
        x (torch.Tensor): The input tensor.

    Returns:
        (torch.Tensor): The output tensor after self-attention.
    """
    B, C, H, W = x.shape
    N = H * W
    qkv = self.qkv(x)
    q, k, v = qkv.view(B, self.num_heads, self.key_dim * 2 + self.head_dim, N).split(
        [self.key_dim, self.key_dim, self.head_dim], dim=2
    )

    attn = (q.transpose(-2, -1) @ k) * self.scale
    attn = attn.softmax(dim=-1)
    x = (v @ attn.transpose(-2, -1)).view(B, C, H, W) + self.pe(v.reshape(B, C, H, W))
    x = self.proj(x)
    return x

ultralytics.nn.modules.block.PSABlock

PSABlock(c, attn_ratio=0.5, num_heads=4, shortcut=True)

Bases: Module

PSABlock class implementing a Position-Sensitive Attention block for neural networks.

This class encapsulates the functionality for applying multi-head attention and feed-forward neural network layers with optional shortcut connections.

Attributes:

Name	Type	Description
`attn`	`Attention`	Multi-head attention module.
`ffn`	`Sequential`	Feed-forward neural network module.
`add`	`bool`	Flag indicating whether to add shortcut connections.

Methods:

Name	Description
`forward`	Performs a forward pass through the PSABlock, applying attention and feed-forward layers.

Examples:

Create a PSABlock and perform a forward pass

>>> psablock = PSABlock(c=128, attn_ratio=0.5, num_heads=4, shortcut=True)
>>> input_tensor = torch.randn(1, 128, 32, 32)
>>> output_tensor = psablock(input_tensor)

Parameters:

Name	Type	Description	Default
`c`	`int`	Input and output channels.	required
`attn_ratio`	`float`	Attention ratio for key dimension.	`0.5`
`num_heads`	`int`	Number of attention heads.	`4`
`shortcut`	`bool`	Whether to use shortcut connections.	`True`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c, attn_ratio=0.5, num_heads=4, shortcut=True) -> None:
    """
    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)
    self.ffn = nn.Sequential(Conv(c, c * 2, 1), Conv(c * 2, c, 1, act=False))
    self.add = shortcut

forward

forward(x)

Execute a forward pass through PSABlock.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor.	required

Returns:

Type	Description
`Tensor`	Output tensor after attention and feed-forward processing.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """
    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

ultralytics.nn.modules.block.PSA

PSA(c1, c2, e=0.5)

Bases: Module

PSA class for implementing Position-Sensitive Attention in neural networks.

This class encapsulates the functionality for applying position-sensitive attention and feed-forward networks to input tensors, enhancing feature extraction and processing capabilities.

Attributes:

Name	Type	Description
`c`	`int`	Number of hidden channels after applying the initial convolution.
`cv1`	`Conv`	1x1 convolution layer to reduce the number of input channels to 2*c.
`cv2`	`Conv`	1x1 convolution layer to reduce the number of output channels to c.
`attn`	`Attention`	Attention module for position-sensitive attention.
`ffn`	`Sequential`	Feed-forward network for further processing.

Methods:

Name	Description
`forward`	Applies position-sensitive attention and feed-forward network to the input tensor.

Examples:

Create a PSA module and apply it to an input tensor

>>> psa = PSA(c1=128, c2=128, e=0.5)
>>> input_tensor = torch.randn(1, 128, 64, 64)
>>> output_tensor = psa.forward(input_tensor)

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`e`	`float`	Expansion ratio.	`0.5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, e=0.5):
    """
    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)
    self.cv1 = Conv(c1, 2 * self.c, 1, 1)
    self.cv2 = Conv(2 * self.c, c1, 1)

    self.attn = Attention(self.c, attn_ratio=0.5, num_heads=self.c // 64)
    self.ffn = nn.Sequential(Conv(self.c, self.c * 2, 1), Conv(self.c * 2, self.c, 1, act=False))

forward

forward(x)

Execute forward pass in PSA module.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor.	required

Returns:

Type	Description
`Tensor`	Output tensor after attention and feed-forward processing.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """
    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)
    return self.cv2(torch.cat((a, b), 1))

ultralytics.nn.modules.block.C2PSA

C2PSA(c1, c2, n=1, e=0.5)

Bases: Module

C2PSA module with attention mechanism for enhanced feature extraction and processing.

This module implements a convolutional block with attention mechanisms to enhance feature extraction and processing capabilities. It includes a series of PSABlock modules for self-attention and feed-forward operations.

Attributes:

Name	Type	Description
`c`	`int`	Number of hidden channels.
`cv1`	`Conv`	1x1 convolution layer to reduce the number of input channels to 2*c.
`cv2`	`Conv`	1x1 convolution layer to reduce the number of output channels to c.
`m`	`Sequential`	Sequential container of PSABlock modules for attention and feed-forward operations.

Methods:

Name	Description
`forward`	Performs a forward pass through the C2PSA module, applying attention and feed-forward operations.

Notes

This module essentially is the same as PSA module, but refactored to allow stacking more PSABlock modules.

Examples:

>>> c2psa = C2PSA(c1=256, c2=256, n=3, e=0.5)
>>> input_tensor = torch.randn(1, 256, 64, 64)
>>> output_tensor = c2psa(input_tensor)

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`n`	`int`	Number of PSABlock modules.	`1`
`e`	`float`	Expansion ratio.	`0.5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=1, e=0.5):
    """
    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)
    self.cv1 = Conv(c1, 2 * self.c, 1, 1)
    self.cv2 = Conv(2 * self.c, c1, 1)

    self.m = nn.Sequential(*(PSABlock(self.c, attn_ratio=0.5, num_heads=self.c // 64) for _ in range(n)))

forward

forward(x)

Process the input tensor through a series of PSA blocks.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor.	required

Returns:

Type	Description
`Tensor`	Output tensor after processing.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """
    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))

ultralytics.nn.modules.block.C2fPSA

C2fPSA(c1, c2, n=1, e=0.5)

Bases: C2f

C2fPSA module with enhanced feature extraction using PSA blocks.

This class extends the C2f module by incorporating PSA blocks for improved attention mechanisms and feature extraction.

Attributes:

Name	Type	Description
`c`	`int`	Number of hidden channels.
`cv1`	`Conv`	1x1 convolution layer to reduce the number of input channels to 2*c.
`cv2`	`Conv`	1x1 convolution layer to reduce the number of output channels to c.
`m`	`ModuleList`	List of PSA blocks for feature extraction.

Methods:

Name	Description
`forward`	Performs a forward pass through the C2fPSA module.
`forward_split`	Performs a forward pass using split() instead of chunk().

Examples:

>>> import torch
>>> from ultralytics.models.common import C2fPSA
>>> model = C2fPSA(c1=64, c2=64, n=3, e=0.5)
>>> x = torch.randn(1, 64, 128, 128)
>>> output = model(x)
>>> print(output.shape)

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`n`	`int`	Number of PSABlock modules.	`1`
`e`	`float`	Expansion ratio.	`0.5`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=1, e=0.5):
    """
    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))

ultralytics.nn.modules.block.SCDown

SCDown(c1, c2, k, s)

Bases: Module

SCDown module for downsampling with separable convolutions.

This module performs downsampling using a combination of pointwise and depthwise convolutions, which helps in efficiently reducing the spatial dimensions of the input tensor while maintaining the channel information.

Attributes:

Name	Type	Description
`cv1`	`Conv`	Pointwise convolution layer that reduces the number of channels.
`cv2`	`Conv`	Depthwise convolution layer that performs spatial downsampling.

Methods:

Name	Description
`forward`	Applies the SCDown module to the input tensor.

Examples:

>>> import torch
>>> from ultralytics import SCDown
>>> model = SCDown(c1=64, c2=128, k=3, s=2)
>>> x = torch.randn(1, 64, 128, 128)
>>> y = model(x)
>>> print(y.shape)
torch.Size([1, 128, 64, 64])

Parameters:

Name	Type	Description	Default
`c1`	`int`	Input channels.	required
`c2`	`int`	Output channels.	required
`k`	`int`	Kernel size.	required
`s`	`int`	Stride.	required

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, k, s):
    """
    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)

forward

forward(x)

Apply convolution and downsampling to the input tensor.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor.	required

Returns:

Type	Description
`Tensor`	Downsampled output tensor.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """
    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))

ultralytics.nn.modules.block.TorchVision

TorchVision(model, weights='DEFAULT', unwrap=True, truncate=2, split=False)

Bases: Module

TorchVision module to allow loading any torchvision model.

This class provides a way to load a model from the torchvision library, optionally load pre-trained weights, and customize the model by truncating or unwrapping layers.

Attributes:

Name	Type	Description
`m`	`Module`	The loaded torchvision model, possibly truncated and unwrapped.

Parameters:

Name	Type	Description	Default
`model`	`str`	Name of the torchvision model to load.	required
`weights`	`str`	Pre-trained weights to load. Default is "DEFAULT".	`'DEFAULT'`
`unwrap`	`bool`	If True, unwraps the model to a sequential containing all but the last `truncate` layers. Default is True.	`True`
`truncate`	`int`	Number of layers to truncate from the end if `unwrap` is True. Default is 2.	`2`
`split`	`bool`	Returns output from intermediate child modules as list. Default is False.	`False`

Parameters:

Name	Type	Description	Default
`model`	`str`	Name of the torchvision model to load.	required
`weights`	`str`	Pre-trained weights to load.	`'DEFAULT'`
`unwrap`	`bool`	Whether to unwrap the model.	`True`
`truncate`	`int`	Number of layers to truncate.	`2`
`split`	`bool`	Whether to split the output.	`False`

Source code in ultralytics/nn/modules/block.py

def __init__(self, model, weights="DEFAULT", unwrap=True, truncate=2, split=False):
    """
    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__()
    if hasattr(torchvision.models, "get_model"):
        self.m = torchvision.models.get_model(model, weights=weights)
    else:
        self.m = torchvision.models.__dict__[model](pretrained=bool(weights))
    if unwrap:
        layers = list(self.m.children())
        if isinstance(layers[0], nn.Sequential):  # Second-level for some models like EfficientNet, Swin
            layers = [*list(layers[0].children()), *layers[1:]]
        self.m = nn.Sequential(*(layers[:-truncate] if truncate else layers))
        self.split = split
    else:
        self.split = False
        self.m.head = self.m.heads = nn.Identity()

forward

forward(x)

Forward pass through the model.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor.	required

Returns:

Type	Description
`Tensor \| List[Tensor]`	Output tensor or list of tensors.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """
    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)
    else:
        y = self.m(x)
    return y

ultralytics.nn.modules.block.AAttn

AAttn(dim, num_heads, area=1)

Bases: Module

Area-attention module for YOLO models, providing efficient attention mechanisms.

This module implements an area-based attention mechanism that processes input features in a spatially-aware manner, making it particularly effective for object detection tasks.

Attributes:

Name	Type	Description
`area`	`int`	Number of areas the feature map is divided.
`num_heads`	`int`	Number of heads into which the attention mechanism is divided.
`head_dim`	`int`	Dimension of each attention head.
`qkv`	`Conv`	Convolution layer for computing query, key and value tensors.
`proj`	`Conv`	Projection convolution layer.
`pe`	`Conv`	Position encoding convolution layer.

Methods:

Name	Description
`forward`	Applies area-attention to input tensor.

Examples:

>>> attn = AAttn(dim=256, num_heads=8, area=4)
>>> x = torch.randn(1, 256, 32, 32)
>>> output = attn(x)
>>> print(output.shape)
torch.Size([1, 256, 32, 32])

Parameters:

Name	Type	Description	Default
`dim`	`int`	Number of hidden channels.	required
`num_heads`	`int`	Number of heads into which the attention mechanism is divided.	required
`area`	`int`	Number of areas the feature map is divided, default is 1.	`1`

Source code in ultralytics/nn/modules/block.py

def __init__(self, dim, num_heads, area=1):
    """
    Initialize an Area-attention module for YOLO models.

    Args:
        dim (int): Number of hidden channels.
        num_heads (int): Number of heads into which the attention mechanism is divided.
        area (int): Number of areas the feature map is divided, default is 1.
    """
    super().__init__()
    self.area = area

    self.num_heads = num_heads
    self.head_dim = head_dim = dim // num_heads
    all_head_dim = head_dim * self.num_heads

    self.qkv = Conv(dim, all_head_dim * 3, 1, act=False)
    self.proj = Conv(all_head_dim, dim, 1, act=False)
    self.pe = Conv(all_head_dim, dim, 7, 1, 3, g=dim, act=False)

forward

forward(x)

Process the input tensor through the area-attention.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor.	required

Returns:

Type	Description
`Tensor`	Output tensor after area-attention.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """
    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

    qkv = self.qkv(x).flatten(2).transpose(1, 2)
    if self.area > 1:
        qkv = qkv.reshape(B * self.area, N // self.area, C * 3)
        B, N, _ = qkv.shape
    q, k, v = (
        qkv.view(B, N, self.num_heads, self.head_dim * 3)
        .permute(0, 2, 3, 1)
        .split([self.head_dim, self.head_dim, self.head_dim], dim=2)
    )
    attn = (q.transpose(-2, -1) @ k) * (self.head_dim**-0.5)
    attn = attn.softmax(dim=-1)
    x = v @ attn.transpose(-2, -1)
    x = x.permute(0, 3, 1, 2)
    v = v.permute(0, 3, 1, 2)

    if self.area > 1:
        x = x.reshape(B // self.area, N * self.area, C)
        v = v.reshape(B // self.area, N * self.area, C)
        B, N, _ = x.shape

    x = x.reshape(B, H, W, C).permute(0, 3, 1, 2).contiguous()
    v = v.reshape(B, H, W, C).permute(0, 3, 1, 2).contiguous()

    x = x + self.pe(v)
    return self.proj(x)

ultralytics.nn.modules.block.ABlock

ABlock(dim, num_heads, mlp_ratio=1.2, area=1)

Bases: Module

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.

Attributes:

Name	Type	Description
`attn`	`AAttn`	Area-attention module for processing spatial features.
`mlp`	`Sequential`	Multi-layer perceptron for feature transformation.

Methods:

Name	Description
`_init_weights`	Initializes module weights using truncated normal distribution.
`forward`	Applies area-attention and feed-forward processing to input tensor.

Examples:

>>> block = ABlock(dim=256, num_heads=8, mlp_ratio=1.2, area=1)
>>> x = torch.randn(1, 256, 32, 32)
>>> output = block(x)
>>> print(output.shape)
torch.Size([1, 256, 32, 32])

Parameters:

Name	Type	Description	Default
`dim`	`int`	Number of input channels.	required
`num_heads`	`int`	Number of heads into which the attention mechanism is divided.	required
`mlp_ratio`	`float`	Expansion ratio for MLP hidden dimension.	`1.2`
`area`	`int`	Number of areas the feature map is divided.	`1`

Source code in ultralytics/nn/modules/block.py

def __init__(self, dim, num_heads, mlp_ratio=1.2, area=1):
    """
    Initialize an Area-attention block module.

    Args:
        dim (int): Number of input channels.
        num_heads (int): Number of heads into which the attention mechanism is divided.
        mlp_ratio (float): Expansion ratio for MLP hidden dimension.
        area (int): Number of areas the feature map is divided.
    """
    super().__init__()

    self.attn = AAttn(dim, num_heads=num_heads, area=area)
    mlp_hidden_dim = int(dim * mlp_ratio)
    self.mlp = nn.Sequential(Conv(dim, mlp_hidden_dim, 1), Conv(mlp_hidden_dim, dim, 1, act=False))

    self.apply(self._init_weights)

forward

forward(x)

Forward pass through ABlock.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor.	required

Returns:

Type	Description
`Tensor`	Output tensor after area-attention and feed-forward processing.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """
    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)

ultralytics.nn.modules.block.A2C2f

A2C2f(
    c1,
    c2,
    n=1,
    a2=True,
    area=1,
    residual=False,
    mlp_ratio=2.0,
    e=0.5,
    g=1,
    shortcut=True,
)

Bases: Module

Area-Attention C2f module for enhanced feature extraction with area-based attention mechanisms.

This module extends the C2f architecture by incorporating area-attention and ABlock layers for improved feature processing. It supports both area-attention and standard convolution modes.

Attributes:

Name	Type	Description
`cv1`	`Conv`	Initial 1x1 convolution layer that reduces input channels to hidden channels.
`cv2`	`Conv`	Final 1x1 convolution layer that processes concatenated features.
`gamma`	`Parameter \| None`	Learnable parameter for residual scaling when using area attention.
`m`	`ModuleList`	List of either ABlock or C3k modules for feature processing.

Methods:

Name	Description
`forward`	Processes input through area-attention or standard convolution pathway.

Examples:

>>> m = A2C2f(512, 512, n=1, a2=True, area=1)
>>> x = torch.randn(1, 512, 32, 32)
>>> output = m(x)
>>> print(output.shape)
torch.Size([1, 512, 32, 32])

Parameters:

Name	Type	Description	Default
`c1`	`int`	Number of input channels.	required
`c2`	`int`	Number of output channels.	required
`n`	`int`	Number of ABlock or C3k modules to stack.	`1`
`a2`	`bool`	Whether to use area attention blocks. If False, uses C3k blocks instead.	`True`
`area`	`int`	Number of areas the feature map is divided.	`1`
`residual`	`bool`	Whether to use residual connections with learnable gamma parameter.	`False`
`mlp_ratio`	`float`	Expansion ratio for MLP hidden dimension.	`2.0`
`e`	`float`	Channel expansion ratio for hidden channels.	`0.5`
`g`	`int`	Number of groups for grouped convolutions.	`1`
`shortcut`	`bool`	Whether to use shortcut connections in C3k blocks.	`True`

Source code in ultralytics/nn/modules/block.py

def __init__(self, c1, c2, n=1, a2=True, area=1, residual=False, mlp_ratio=2.0, e=0.5, g=1, shortcut=True):
    """
    Initialize Area-Attention C2f module.

    Args:
        c1 (int): Number of input channels.
        c2 (int): Number of output channels.
        n (int): Number of ABlock or C3k modules to stack.
        a2 (bool): Whether to use area attention blocks. If False, uses C3k blocks instead.
        area (int): Number of areas the feature map is divided.
        residual (bool): Whether to use residual connections with learnable gamma parameter.
        mlp_ratio (float): Expansion ratio for MLP hidden dimension.
        e (float): Channel expansion ratio for hidden channels.
        g (int): Number of groups for grouped convolutions.
        shortcut (bool): Whether to use shortcut connections in C3k blocks.
    """
    super().__init__()
    c_ = int(c2 * e)  # hidden channels
    assert c_ % 32 == 0, "Dimension of ABlock be a multiple of 32."

    self.cv1 = Conv(c1, c_, 1, 1)
    self.cv2 = Conv((1 + n) * c_, c2, 1)

    self.gamma = nn.Parameter(0.01 * torch.ones(c2), requires_grad=True) if a2 and residual else None
    self.m = nn.ModuleList(
        nn.Sequential(*(ABlock(c_, c_ // 32, mlp_ratio, area) for _ in range(2)))
        if a2
        else C3k(c_, c_, 2, shortcut, g)
        for _ in range(n)
    )

forward

forward(x)

Forward pass through A2C2f layer.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor.	required

Returns:

Type	Description
`Tensor`	Output tensor after processing.

Source code in ultralytics/nn/modules/block.py

def forward(self, x):
    """
    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))
    if self.gamma is not None:
        return x + self.gamma.view(-1, len(self.gamma), 1, 1) * y
    return y

📅 Created 1 year ago ✏️ Updated 28 days ago

Reference for ultralytics/nn/modules/block.py

ultralytics.nn.modules.block.DFL

forward

ultralytics.nn.modules.block.Proto

forward

ultralytics.nn.modules.block.HGStem

forward

ultralytics.nn.modules.block.HGBlock

forward

ultralytics.nn.modules.block.SPP

forward

ultralytics.nn.modules.block.SPPF

forward

ultralytics.nn.modules.block.C1

forward

ultralytics.nn.modules.block.C2

forward

ultralytics.nn.modules.block.C2f

forward

forward_split

ultralytics.nn.modules.block.C3

forward

ultralytics.nn.modules.block.C3x

ultralytics.nn.modules.block.RepC3

forward

ultralytics.nn.modules.block.C3TR

ultralytics.nn.modules.block.C3Ghost

ultralytics.nn.modules.block.GhostBottleneck

forward

ultralytics.nn.modules.block.Bottleneck

forward

ultralytics.nn.modules.block.BottleneckCSP

forward

ultralytics.nn.modules.block.ResNetBlock

forward

ultralytics.nn.modules.block.ResNetLayer

forward

ultralytics.nn.modules.block.MaxSigmoidAttnBlock

forward

ultralytics.nn.modules.block.C2fAttn

forward

forward_split

ultralytics.nn.modules.block.ImagePoolingAttn

forward

ultralytics.nn.modules.block.ContrastiveHead

forward

ultralytics.nn.modules.block.BNContrastiveHead

forward

ultralytics.nn.modules.block.RepBottleneck

ultralytics.nn.modules.block.RepCSP

ultralytics.nn.modules.block.RepNCSPELAN4

forward

forward_split

ultralytics.nn.modules.block.ELAN1

ultralytics.nn.modules.block.AConv

forward

ultralytics.nn.modules.block.ADown

forward

ultralytics.nn.modules.block.SPPELAN

forward

ultralytics.nn.modules.block.CBLinear

forward

ultralytics.nn.modules.block.CBFuse

forward

ultralytics.nn.modules.block.C3f

forward

ultralytics.nn.modules.block.C3k2

ultralytics.nn.modules.block.C3k

ultralytics.nn.modules.block.RepVGGDW

forward

forward_fuse

fuse

ultralytics.nn.modules.block.CIB

forward

ultralytics.nn.modules.block.C2fCIB

ultralytics.nn.modules.block.Attention

forward

ultralytics.nn.modules.block.PSABlock

forward

ultralytics.nn.modules.block.PSA

Reference for `ultralytics/nn/modules/block.py`