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
246
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.

Source code in ultralytics/nn/modules/block.py 
252
253
254
255
256
257
258
259
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 
261
262
263
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 
269
270
271
272
273
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 
279
280
281
282
283
284
285
286
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, 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 RT-DETR neck layer.

Source code in ultralytics/nn/modules/block.py 
288
289
290
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 
296
297
298
299
300
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 
306
307
308
309
310
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 
316
317
318
319
320
321
322
323
324
325
326
327
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 
329
330
331
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 
337
338
339
340
341
342
343
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 
345
346
347
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 
353
354
355
356
357
358
359
360
361
362
363
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 
365
366
367
368
369
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 
375
376
377
378
379
380
381
382
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 
384
385
386
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 
392
393
394
395
396
397
398
399
400
401
402
403
404
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 
406
407
408
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 
414
415
416
417
418
419
420
421
422
423
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 
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
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 
449
450
451
452
453
454
455
456
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 
458
459
460
461
462
463
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 
465
466
467
468
469
470
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 
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
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 
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
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 
522
523
524
525
526
527
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 
529
530
531
532
533
534
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:

Name Type Description Default
embed_dims int

Embed dimensions of text and image features.

 required
Source code in ultralytics/nn/modules/block.py 
545
546
547
548
549
550
551
552
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 
554
555
556
557
558
559
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 
565
566
567
568
569
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 
575
576
577
578
579
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 
585
586
587
588
589
590
591
592
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 
594
595
596
597
598
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 
600
601
602
603
604
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 
610
611
612
613
614
615
616
617
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 
623
624
625
626
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 
628
629
630
631
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 
637
638
639
640
641
642
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 
644
645
646
647
648
649
650
651
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 
657
658
659
660
661
662
663
664
665
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 
667
668
669
670
671
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 
677
678
679
680
681
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 
683
684
685
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 
691
692
693
694
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 
696
697
698
699
700
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 
706
707
708
709
710
711
712
713
714
715
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 
717
718
719
720
721
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 
727
728
729
730
731
732
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 
738
739
740
741
742
743
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 
749
750
751
752
753
754
755
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:

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 
757
758
759
760
761
762
763
764
765
766
767
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:

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 
769
770
771
772
773
774
775
776
777
778
779
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 
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
@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:

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
Source code in ultralytics/nn/modules/block.py 
820
821
822
823
824
825
826
827
828
829
830
831
832
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:

Name Type Description Default
x Tensor

Input tensor.

 required

Returns:

Type Description
Tensor

Output tensor.
Source code in ultralytics/nn/modules/block.py 
834
835
836
837
838
839
840
841
842
843
844
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
Source code in ultralytics/nn/modules/block.py 
861
862
863
864
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:

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

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 
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
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)
Source code in ultralytics/nn/modules/block.py 
945
946
947
948
949
950
951
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 
953
954
955
956
957
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:

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)
Source code in ultralytics/nn/modules/block.py 
984
985
986
987
988
989
990
991
992
993
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 
 995
 996
 997
 998
 999
1000
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:

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)
Source code in ultralytics/nn/modules/block.py 
1028
1029
1030
1031
1032
1033
1034
1035
1036
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 
1038
1039
1040
1041
1042
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:

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

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])
Source code in ultralytics/nn/modules/block.py 
1101
1102
1103
1104
1105
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 
1107
1108
1109
def forward(self, x):
    """Applies convolution and downsampling to the input tensor in the SCDown module."""
    return self.cv2(self.cv1(x))



📅 Created 1 year ago ✏️ Updated 2 months ago
glenn-jocher jk4e Burhan-Q Laughing-q