Skip to content

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 
62
63
64
65
66
67
68
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)

Applies a transformer layer on input tensor 'x' and returns a tensor.

Source code in ultralytics/nn/modules/block.py 
70
71
72
73
def forward(self, x):
    """Applies a transformer layer on input tensor 'x' and returns a tensor."""
    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.

Input arguments are ch_in, number of protos, number of masks.

Source code in ultralytics/nn/modules/block.py 
80
81
82
83
84
85
86
87
88
89
90
def __init__(self, c1, c_=256, c2=32):
    """
    Initializes the YOLOv8 mask Proto module with specified number of protos and masks.

    Input arguments are ch_in, number of protos, number of masks.
    """
    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)

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

Source code in ultralytics/nn/modules/block.py 
92
93
94
def forward(self, x):
    """Performs 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

Source code in ultralytics/nn/modules/block.py 
104
105
106
107
108
109
110
111
112
def __init__(self, c1, cm, c2):
    """Initialize the SPP layer with input/output channels and specified kernel sizes for max pooling."""
    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 
114
115
116
117
118
119
120
121
122
123
124
125
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

Source code in ultralytics/nn/modules/block.py 
135
136
137
138
139
140
141
142
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."""
    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 
144
145
146
147
148
149
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.

Source code in ultralytics/nn/modules/block.py 
155
156
157
158
159
160
161
def __init__(self, c1, c2, k=(5, 9, 13)):
    """Initialize the SPP layer with input/output channels and pooling kernel sizes."""
    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 
163
164
165
166
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.

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

Source code in ultralytics/nn/modules/block.py 
172
173
174
175
176
177
178
179
180
181
182
def __init__(self, c1, c2, k=5):
    """
    Initializes the SPPF layer with given input/output channels and kernel size.

    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)

Forward pass through Ghost Convolution block.

Source code in ultralytics/nn/modules/block.py 
184
185
186
187
188
def forward(self, x):
    """Forward pass through Ghost Convolution block."""
    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.

Source code in ultralytics/nn/modules/block.py 
194
195
196
197
198
def __init__(self, c1, c2, n=1):
    """Initializes the CSP Bottleneck with configurations for 1 convolution with arguments ch_in, ch_out, number."""
    super().__init__()
    self.cv1 = Conv(c1, c2, 1, 1)
    self.m = nn.Sequential(*(Conv(c2, c2, 3) for _ in range(n)))

forward

forward(x)

Applies cross-convolutions to input in the C3 module.

Source code in ultralytics/nn/modules/block.py 
200
201
202
203
def forward(self, x):
    """Applies cross-convolutions to input in the C3 module."""
    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.

Source code in ultralytics/nn/modules/block.py 
209
210
211
212
213
214
215
216
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."""
    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 
218
219
220
221
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.

Source code in ultralytics/nn/modules/block.py 
227
228
229
230
231
232
233
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."""
    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 
235
236
237
238
239
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 
241
242
243
244
245
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.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.

Source code in ultralytics/nn/modules/block.py 
251
252
253
254
255
256
257
258
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."""
    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 2 convolutions.

Source code in ultralytics/nn/modules/block.py 
260
261
262
def forward(self, x):
    """Forward pass through the CSP bottleneck with 2 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.

Source code in ultralytics/nn/modules/block.py 
268
269
270
271
272
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
    """Initialize C3TR instance and set default parameters."""
    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.

Source code in ultralytics/nn/modules/block.py 
278
279
280
281
282
283
284
285
def __init__(self, c1, c2, n=3, e=1.0):
    """Initialize CSP Bottleneck with a single convolution using input channels, output channels, and number."""
    super().__init__()
    c_ = int(c2 * e)  # hidden channels
    self.cv1 = Conv(c1, c2, 1, 1)
    self.cv2 = Conv(c1, c2, 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 RT-DETR neck layer.

Source code in ultralytics/nn/modules/block.py 
287
288
289
def forward(self, x):
    """Forward pass of RT-DETR neck layer."""
    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().

Source code in ultralytics/nn/modules/block.py 
295
296
297
298
299
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
    """Initialize C3Ghost module with GhostBottleneck()."""
    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().

Source code in ultralytics/nn/modules/block.py 
305
306
307
308
309
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."""
    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.

Source code in ultralytics/nn/modules/block.py 
315
316
317
318
319
320
321
322
323
324
325
326
def __init__(self, c1, c2, k=3, s=1):
    """Initializes GhostBottleneck module with arguments ch_in, ch_out, kernel, 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)

Applies skip connection and concatenation to input tensor.

Source code in ultralytics/nn/modules/block.py 
328
329
330
def forward(self, x):
    """Applies 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.

Source code in ultralytics/nn/modules/block.py 
336
337
338
339
340
341
342
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."""
    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)

Applies the YOLO FPN to input data.

Source code in ultralytics/nn/modules/block.py 
344
345
346
def forward(self, x):
    """Applies the YOLO FPN to input data."""
    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.

Source code in ultralytics/nn/modules/block.py 
352
353
354
355
356
357
358
359
360
361
362
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."""
    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)

Applies a CSP bottleneck with 3 convolutions.

Source code in ultralytics/nn/modules/block.py 
364
365
366
367
368
def forward(self, x):
    """Applies a 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.

Source code in ultralytics/nn/modules/block.py 
374
375
376
377
378
379
380
381
def __init__(self, c1, c2, s=1, e=4):
    """Initialize convolution with given parameters."""
    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 
383
384
385
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.

Source code in ultralytics/nn/modules/block.py 
391
392
393
394
395
396
397
398
399
400
401
402
403
def __init__(self, c1, c2, s=1, is_first=False, n=1, e=4):
    """Initializes the ResNetLayer given arguments."""
    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 
405
406
407
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.

Source code in ultralytics/nn/modules/block.py 
413
414
415
416
417
418
419
420
421
422
def __init__(self, c1, c2, nh=1, ec=128, gc=512, scale=False):
    """Initializes MaxSigmoidAttnBlock with specified arguments."""
    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 process.

Source code in ultralytics/nn/modules/block.py 
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
def forward(self, x, guide):
    """Forward process."""
    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.

Source code in ultralytics/nn/modules/block.py 
448
449
450
451
452
453
454
455
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."""
    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.

Source code in ultralytics/nn/modules/block.py 
457
458
459
460
461
462
def forward(self, x, guide):
    """Forward pass through C2f layer."""
    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().

Source code in ultralytics/nn/modules/block.py 
464
465
466
467
468
469
def forward_split(self, x, guide):
    """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.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.

Source code in ultralytics/nn/modules/block.py 
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
def __init__(self, ec=256, ch=(), ct=512, nh=8, k=3, scale=False):
    """Initializes ImagePoolingAttn with specified arguments."""
    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)

Executes attention mechanism on input tensor x and guide tensor.

Source code in ultralytics/nn/modules/block.py 
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
def forward(self, x, text):
    """Executes attention mechanism on input tensor x and guide tensor."""
    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 
521
522
523
524
525
526
def __init__(self):
    """Initializes ContrastiveHead with specified 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.

Source code in ultralytics/nn/modules/block.py 
528
529
530
531
532
533
def forward(self, x, w):
    """Forward function of contrastive learning."""
    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:

NameTypeDescriptionDefault
embed_dimsint

Embed dimensions of text and image features.

required
Source code in ultralytics/nn/modules/block.py 
544
545
546
547
548
549
550
551
def __init__(self, embed_dims: int):
    """Initialize ContrastiveHead with region-text similarity parameters."""
    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.

Source code in ultralytics/nn/modules/block.py 
553
554
555
556
557
558
def forward(self, x, w):
    """Forward function of contrastive learning."""
    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.

Source code in ultralytics/nn/modules/block.py 
564
565
566
567
568
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."""
    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.

Source code in ultralytics/nn/modules/block.py 
574
575
576
577
578
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."""
    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.

Source code in ultralytics/nn/modules/block.py 
584
585
586
587
588
589
590
591
def __init__(self, c1, c2, c3, c4, n=1):
    """Initializes CSP-ELAN layer with specified channel sizes, repetitions, and convolutions."""
    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 
593
594
595
596
597
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 
599
600
601
602
603
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.

Source code in ultralytics/nn/modules/block.py 
609
610
611
612
613
614
615
616
def __init__(self, c1, c2, c3, c4):
    """Initializes ELAN1 layer with specified channel sizes."""
    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.

Source code in ultralytics/nn/modules/block.py 
622
623
624
625
def __init__(self, c1, c2):
    """Initializes AConv module with convolution layers."""
    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 
627
628
629
630
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.

Source code in ultralytics/nn/modules/block.py 
636
637
638
639
640
641
def __init__(self, c1, c2):
    """Initializes ADown module with convolution layers to downsample input from channels c1 to c2."""
    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 
643
644
645
646
647
648
649
650
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.

Source code in ultralytics/nn/modules/block.py 
656
657
658
659
660
661
662
663
664
def __init__(self, c1, c2, c3, k=5):
    """Initializes SPP-ELAN block with convolution and max pooling layers for spatial pyramid 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 
666
667
668
669
670
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.

Source code in ultralytics/nn/modules/block.py 
676
677
678
679
680
def __init__(self, c1, c2s, k=1, s=1, p=None, g=1):
    """Initializes the CBLinear module, passing inputs unchanged."""
    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 
682
683
684
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.

Source code in ultralytics/nn/modules/block.py 
690
691
692
693
def __init__(self, idx):
    """Initializes CBFuse module with layer index for selective feature fusion."""
    super().__init__()
    self.idx = idx

forward

forward(xs)

Forward pass through CBFuse layer.

Source code in ultralytics/nn/modules/block.py 
695
696
697
698
699
def forward(self, xs):
    """Forward pass through CBFuse layer."""
    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.

Source code in ultralytics/nn/modules/block.py 
705
706
707
708
709
710
711
712
713
714
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.
    """
    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 C2f layer.

Source code in ultralytics/nn/modules/block.py 
716
717
718
719
720
def forward(self, x):
    """Forward pass through C2f 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.

Source code in ultralytics/nn/modules/block.py 
726
727
728
729
730
731
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."""
    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.

Source code in ultralytics/nn/modules/block.py 
737
738
739
740
741
742
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."""
    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.

Source code in ultralytics/nn/modules/block.py 
748
749
750
751
752
753
754
def __init__(self, ed) -> None:
    """Initializes RepVGGDW with depthwise separable convolutional layers for efficient processing."""
    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)

Performs a forward pass of the RepVGGDW block.

Parameters:

NameTypeDescriptionDefault
xTensor

Input tensor.

required

Returns:

TypeDescription
Tensor

Output tensor after applying the depth wise separable convolution.

Source code in ultralytics/nn/modules/block.py 
756
757
758
759
760
761
762
763
764
765
766
def forward(self, x):
    """
    Performs 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)

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

Parameters:

NameTypeDescriptionDefault
xTensor

Input tensor.

required

Returns:

TypeDescription
Tensor

Output tensor after applying the depth wise separable convolution.

Source code in ultralytics/nn/modules/block.py 
768
769
770
771
772
773
774
775
776
777
778
def forward_fuse(self, x):
    """
    Performs 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()

Fuses 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 
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
@torch.no_grad()
def fuse(self):
    """
    Fuses 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:

NameTypeDescriptionDefault
c1int

Number of input channels.

required
c2int

Number of output channels.

required
shortcutbool

Whether to add a shortcut connection. Defaults to True.

True
efloat

Scaling factor for the hidden channels. Defaults to 0.5.

0.5
lkbool

Whether to use RepVGGDW for the third convolutional layer. Defaults to False.

False
Source code in ultralytics/nn/modules/block.py 
819
820
821
822
823
824
825
826
827
828
829
830
831
def __init__(self, c1, c2, shortcut=True, e=0.5, lk=False):
    """Initializes the custom model with optional shortcut, scaling factor, and RepVGGDW layer."""
    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:

NameTypeDescriptionDefault
xTensor

Input tensor.

required

Returns:

TypeDescription
Tensor

Output tensor.

Source code in ultralytics/nn/modules/block.py 
833
834
835
836
837
838
839
840
841
842
843
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:

NameTypeDescriptionDefault
c1int

Number of input channels.

required
c2int

Number of output channels.

required
nint

Number of CIB modules to stack. Defaults to 1.

1
shortcutbool

Whether to use shortcut connection. Defaults to False.

False
lkbool

Whether to use local key connection. Defaults to False.

False
gint

Number of groups for grouped convolution. Defaults to 1.

1
efloat

Expansion ratio for CIB modules. Defaults to 0.5.

0.5
Source code in ultralytics/nn/modules/block.py 
860
861
862
863
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."""
    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:

NameTypeDescriptionDefault
dimint

The input tensor dimension.

required
num_headsint

The number of attention heads.

8
attn_ratiofloat

The ratio of the attention key dimension to the head dimension.

0.5

Attributes:

NameTypeDescription
num_headsint

The number of attention heads.

head_dimint

The dimension of each attention head.

key_dimint

The dimension of the attention key.

scalefloat

The scaling factor for the attention scores.

qkvConv

Convolutional layer for computing the query, key, and value.

projConv

Convolutional layer for projecting the attended values.

peConv

Convolutional layer for positional encoding.

Source code in ultralytics/nn/modules/block.py 
885
886
887
888
889
890
891
892
893
894
895
896
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."""
    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:

NameTypeDescriptionDefault
xTensor

The input tensor.

required

Returns:

TypeDescription
Tensor

The output tensor after self-attention.

Source code in ultralytics/nn/modules/block.py 
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
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:

NameTypeDescription
attnAttention

Multi-head attention module.

ffnSequential

Feed-forward neural network module.

addbool

Flag indicating whether to add shortcut connections.

Methods:

NameDescription
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)
Source code in ultralytics/nn/modules/block.py 
944
945
946
947
948
949
950
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."""
    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)

Executes a forward pass through PSABlock, applying attention and feed-forward layers to the input tensor.

Source code in ultralytics/nn/modules/block.py 
952
953
954
955
956
def forward(self, x):
    """Executes a forward pass through PSABlock, applying attention and feed-forward layers to the input tensor."""
    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:

NameTypeDescription
cint

Number of hidden channels after applying the initial convolution.

cv1Conv

1x1 convolution layer to reduce the number of input channels to 2*c.

cv2Conv

1x1 convolution layer to reduce the number of output channels to c.

attnAttention

Attention module for position-sensitive attention.

ffnSequential

Feed-forward network for further processing.

Methods:

NameDescription
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)
Source code in ultralytics/nn/modules/block.py 
983
984
985
986
987
988
989
990
991
992
def __init__(self, c1, c2, e=0.5):
    """Initializes the PSA module with input/output channels and attention mechanism for feature extraction."""
    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)

Executes forward pass in PSA module, applying attention and feed-forward layers to the input tensor.

Source code in ultralytics/nn/modules/block.py 
994
995
996
997
998
999
def forward(self, x):
    """Executes forward pass in PSA module, applying attention and feed-forward layers to the input tensor."""
    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:

NameTypeDescription
cint

Number of hidden channels.

cv1Conv

1x1 convolution layer to reduce the number of input channels to 2*c.

cv2Conv

1x1 convolution layer to reduce the number of output channels to c.

mSequential

Sequential container of PSABlock modules for attention and feed-forward operations.

Methods:

NameDescription
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)
Source code in ultralytics/nn/modules/block.py 
1027
1028
1029
1030
1031
1032
1033
1034
1035
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."""
    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)

Processes the input tensor 'x' through a series of PSA blocks and returns the transformed tensor.

Source code in ultralytics/nn/modules/block.py 
1037
1038
1039
1040
1041
def forward(self, x):
    """Processes the input tensor 'x' through a series of PSA blocks and returns the transformed tensor."""
    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:

NameTypeDescription
cint

Number of hidden channels.

cv1Conv

1x1 convolution layer to reduce the number of input channels to 2*c.

cv2Conv

1x1 convolution layer to reduce the number of output channels to c.

mModuleList

List of PSA blocks for feature extraction.

Methods:

NameDescription
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)
Source code in ultralytics/nn/modules/block.py 
1069
1070
1071
1072
1073
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."""
    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:

NameTypeDescription
cv1Conv

Pointwise convolution layer that reduces the number of channels.

cv2Conv

Depthwise convolution layer that performs spatial downsampling.

Methods:

NameDescription
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])
Source code in ultralytics/nn/modules/block.py 
1100
1101
1102
1103
1104
def __init__(self, c1, c2, k, s):
    """Initializes the SCDown module with specified input/output channels, kernel size, and 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)

Applies convolution and downsampling to the input tensor in the SCDown module.

Source code in ultralytics/nn/modules/block.py 
1106
1107
1108
def forward(self, x):
    """Applies convolution and downsampling to the input tensor in the SCDown module."""
    return self.cv2(self.cv1(x))



📅 Created 11 months ago ✏️ Updated 21 days ago
glenn-jocher jk4e Burhan-Q Laughing-q