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
63
64
65
66
67
68
69
def __init__(self, c1=16):
    """Initialize a convolutional layer with a given number of input channels."""
    super().__init__()
    self.conv = nn.Conv2d(c1, 1, 1, bias=False).requires_grad_(False)
    x = torch.arange(c1, dtype=torch.float)
    self.conv.weight.data[:] = nn.Parameter(x.view(1, c1, 1, 1))
    self.c1 = c1

forward

forward(x)

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

Source code in ultralytics/nn/modules/block.py
71
72
73
74
def forward(self, x):
    """Apply the DFL module to input tensor and return transformed output."""
    b, _, a = x.shape  # batch, channels, anchors
    return self.conv(x.view(b, 4, self.c1, a).transpose(2, 1).softmax(1)).view(b, 4, a)





ultralytics.nn.modules.block.Proto

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

Bases: Module

YOLOv8 mask Proto module for segmentation models.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c_ int

Intermediate channels.

256
c2 int

Output channels (number of protos).

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

    Args:
        c1 (int): Input channels.
        c_ (int): Intermediate channels.
        c2 (int): Output channels (number of protos).
    """
    super().__init__()
    self.cv1 = Conv(c1, c_, k=3)
    self.upsample = nn.ConvTranspose2d(c_, c_, 2, 2, 0, bias=True)  # nn.Upsample(scale_factor=2, mode='nearest')
    self.cv2 = Conv(c_, c_, k=3)
    self.cv3 = Conv(c_, c2)

forward

forward(x)

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

Source code in ultralytics/nn/modules/block.py
96
97
98
def forward(self, x):
    """Perform a forward pass through layers using an upsampled input image."""
    return self.cv3(self.cv2(self.upsample(self.cv1(x))))





ultralytics.nn.modules.block.HGStem

HGStem(c1, cm, c2)

Bases: Module

StemBlock of PPHGNetV2 with 5 convolutions and one maxpool2d.

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

Parameters:

Name Type Description Default
c1 int

Input channels.

required
cm int

Middle channels.

required
c2 int

Output channels.

required
Source code in ultralytics/nn/modules/block.py
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def __init__(self, c1, cm, c2):
    """
    Initialize the StemBlock of PPHGNetV2.

    Args:
        c1 (int): Input channels.
        cm (int): Middle channels.
        c2 (int): Output channels.
    """
    super().__init__()
    self.stem1 = Conv(c1, cm, 3, 2, act=nn.ReLU())
    self.stem2a = Conv(cm, cm // 2, 2, 1, 0, act=nn.ReLU())
    self.stem2b = Conv(cm // 2, cm, 2, 1, 0, act=nn.ReLU())
    self.stem3 = Conv(cm * 2, cm, 3, 2, act=nn.ReLU())
    self.stem4 = Conv(cm, c2, 1, 1, act=nn.ReLU())
    self.pool = nn.MaxPool2d(kernel_size=2, stride=1, padding=0, ceil_mode=True)

forward

forward(x)

Forward pass of a PPHGNetV2 backbone layer.

Source code in ultralytics/nn/modules/block.py
125
126
127
128
129
130
131
132
133
134
135
136
def forward(self, x):
    """Forward pass of a PPHGNetV2 backbone layer."""
    x = self.stem1(x)
    x = F.pad(x, [0, 1, 0, 1])
    x2 = self.stem2a(x)
    x2 = F.pad(x2, [0, 1, 0, 1])
    x2 = self.stem2b(x2)
    x1 = self.pool(x)
    x = torch.cat([x1, x2], dim=1)
    x = self.stem3(x)
    x = self.stem4(x)
    return x





ultralytics.nn.modules.block.HGBlock

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

Bases: Module

HG_Block of PPHGNetV2 with 2 convolutions and LightConv.

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

Parameters:

Name Type Description Default
c1 int

Input channels.

required
cm int

Middle channels.

required
c2 int

Output channels.

required
k int

Kernel size.

3
n int

Number of LightConv or Conv blocks.

6
lightconv bool

Whether to use LightConv.

False
shortcut bool

Whether to use shortcut connection.

False
act Module

Activation function.

ReLU()
Source code in ultralytics/nn/modules/block.py
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
def __init__(self, c1, cm, c2, k=3, n=6, lightconv=False, shortcut=False, act=nn.ReLU()):
    """
    Initialize HGBlock with specified parameters.

    Args:
        c1 (int): Input channels.
        cm (int): Middle channels.
        c2 (int): Output channels.
        k (int): Kernel size.
        n (int): Number of LightConv or Conv blocks.
        lightconv (bool): Whether to use LightConv.
        shortcut (bool): Whether to use shortcut connection.
        act (nn.Module): Activation function.
    """
    super().__init__()
    block = LightConv if lightconv else Conv
    self.m = nn.ModuleList(block(c1 if i == 0 else cm, cm, k=k, act=act) for i in range(n))
    self.sc = Conv(c1 + n * cm, c2 // 2, 1, 1, act=act)  # squeeze conv
    self.ec = Conv(c2 // 2, c2, 1, 1, act=act)  # excitation conv
    self.add = shortcut and c1 == c2

forward

forward(x)

Forward pass of a PPHGNetV2 backbone layer.

Source code in ultralytics/nn/modules/block.py
167
168
169
170
171
172
def forward(self, x):
    """Forward pass of a PPHGNetV2 backbone layer."""
    y = [x]
    y.extend(m(y[-1]) for m in self.m)
    y = self.ec(self.sc(torch.cat(y, 1)))
    return y + x if self.add else y





ultralytics.nn.modules.block.SPP

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

Bases: Module

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

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
k Tuple[int, int, int]

Kernel sizes for max pooling.

(5, 9, 13)
Source code in ultralytics/nn/modules/block.py
178
179
180
181
182
183
184
185
186
187
188
189
190
191
def __init__(self, c1, c2, k=(5, 9, 13)):
    """
    Initialize the SPP layer with input/output channels and pooling kernel sizes.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        k (Tuple[int, int, int]): Kernel sizes for max pooling.
    """
    super().__init__()
    c_ = c1 // 2  # hidden channels
    self.cv1 = Conv(c1, c_, 1, 1)
    self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
    self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])

forward

forward(x)

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

Source code in ultralytics/nn/modules/block.py
193
194
195
196
def forward(self, x):
    """Forward pass of the SPP layer, performing spatial pyramid pooling."""
    x = self.cv1(x)
    return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))





ultralytics.nn.modules.block.SPPF

SPPF(c1, c2, k=5)

Bases: Module

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

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
k int

Kernel size.

5
Notes

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

Source code in ultralytics/nn/modules/block.py
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
def __init__(self, c1, c2, k=5):
    """
    Initialize the SPPF layer with given input/output channels and kernel size.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        k (int): Kernel size.

    Notes:
        This module is equivalent to SPP(k=(5, 9, 13)).
    """
    super().__init__()
    c_ = c1 // 2  # hidden channels
    self.cv1 = Conv(c1, c_, 1, 1)
    self.cv2 = Conv(c_ * 4, c2, 1, 1)
    self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)

forward

forward(x)

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

Source code in ultralytics/nn/modules/block.py
220
221
222
223
224
def forward(self, x):
    """Apply sequential pooling operations to input and return concatenated feature maps."""
    y = [self.cv1(x)]
    y.extend(self.m(y[-1]) for _ in range(3))
    return self.cv2(torch.cat(y, 1))





ultralytics.nn.modules.block.C1

C1(c1, c2, n=1)

Bases: Module

CSP Bottleneck with 1 convolution.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
n int

Number of convolutions.

1
Source code in ultralytics/nn/modules/block.py
230
231
232
233
234
235
236
237
238
239
240
241
def __init__(self, c1, c2, n=1):
    """
    Initialize the CSP Bottleneck with 1 convolution.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        n (int): Number of convolutions.
    """
    super().__init__()
    self.cv1 = Conv(c1, c2, 1, 1)
    self.m = nn.Sequential(*(Conv(c2, c2, 3) for _ in range(n)))

forward

forward(x)

Apply convolution and residual connection to input tensor.

Source code in ultralytics/nn/modules/block.py
243
244
245
246
def forward(self, x):
    """Apply convolution and residual connection to input tensor."""
    y = self.cv1(x)
    return self.m(y) + y





ultralytics.nn.modules.block.C2

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

Bases: Module

CSP Bottleneck with 2 convolutions.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
n int

Number of Bottleneck blocks.

1
shortcut bool

Whether to use shortcut connections.

True
g int

Groups for convolutions.

1
e float

Expansion ratio.

0.5
Source code in ultralytics/nn/modules/block.py
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
    """
    Initialize a CSP Bottleneck with 2 convolutions.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        n (int): Number of Bottleneck blocks.
        shortcut (bool): Whether to use shortcut connections.
        g (int): Groups for convolutions.
        e (float): Expansion ratio.
    """
    super().__init__()
    self.c = int(c2 * e)  # hidden channels
    self.cv1 = Conv(c1, 2 * self.c, 1, 1)
    self.cv2 = Conv(2 * self.c, c2, 1)  # optional act=FReLU(c2)
    # self.attention = ChannelAttention(2 * self.c)  # or SpatialAttention()
    self.m = nn.Sequential(*(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n)))

forward

forward(x)

Forward pass through the CSP bottleneck with 2 convolutions.

Source code in ultralytics/nn/modules/block.py
271
272
273
274
def forward(self, x):
    """Forward pass through the CSP bottleneck with 2 convolutions."""
    a, b = self.cv1(x).chunk(2, 1)
    return self.cv2(torch.cat((self.m(a), b), 1))





ultralytics.nn.modules.block.C2f

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

Bases: Module

Faster Implementation of CSP Bottleneck with 2 convolutions.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
n int

Number of Bottleneck blocks.

1
shortcut bool

Whether to use shortcut connections.

False
g int

Groups for convolutions.

1
e float

Expansion ratio.

0.5
Source code in ultralytics/nn/modules/block.py
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
    """
    Initialize a CSP bottleneck with 2 convolutions.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        n (int): Number of Bottleneck blocks.
        shortcut (bool): Whether to use shortcut connections.
        g (int): Groups for convolutions.
        e (float): Expansion ratio.
    """
    super().__init__()
    self.c = int(c2 * e)  # hidden channels
    self.cv1 = Conv(c1, 2 * self.c, 1, 1)
    self.cv2 = Conv((2 + n) * self.c, c2, 1)  # optional act=FReLU(c2)
    self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))

forward

forward(x)

Forward pass through C2f layer.

Source code in ultralytics/nn/modules/block.py
298
299
300
301
302
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
304
305
306
307
308
309
def forward_split(self, x):
    """Forward pass using split() instead of chunk()."""
    y = self.cv1(x).split((self.c, self.c), 1)
    y = [y[0], y[1]]
    y.extend(m(y[-1]) for m in self.m)
    return self.cv2(torch.cat(y, 1))





ultralytics.nn.modules.block.C3

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

Bases: Module

CSP Bottleneck with 3 convolutions.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
n int

Number of Bottleneck blocks.

1
shortcut bool

Whether to use shortcut connections.

True
g int

Groups for convolutions.

1
e float

Expansion ratio.

0.5
Source code in ultralytics/nn/modules/block.py
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
    """
    Initialize the CSP Bottleneck with 3 convolutions.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        n (int): Number of Bottleneck blocks.
        shortcut (bool): Whether to use shortcut connections.
        g (int): Groups for convolutions.
        e (float): Expansion ratio.
    """
    super().__init__()
    c_ = int(c2 * e)  # hidden channels
    self.cv1 = Conv(c1, c_, 1, 1)
    self.cv2 = Conv(c1, c_, 1, 1)
    self.cv3 = Conv(2 * c_, c2, 1)  # optional act=FReLU(c2)
    self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, k=((1, 1), (3, 3)), e=1.0) for _ in range(n)))

forward

forward(x)

Forward pass through the CSP bottleneck with 3 convolutions.

Source code in ultralytics/nn/modules/block.py
334
335
336
def forward(self, x):
    """Forward pass through the CSP bottleneck with 3 convolutions."""
    return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))





ultralytics.nn.modules.block.C3x

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

Bases: C3

C3 module with cross-convolutions.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
n int

Number of Bottleneck blocks.

1
shortcut bool

Whether to use shortcut connections.

True
g int

Groups for convolutions.

1
e float

Expansion ratio.

0.5
Source code in ultralytics/nn/modules/block.py
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
    """
    Initialize C3 module with cross-convolutions.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        n (int): Number of Bottleneck blocks.
        shortcut (bool): Whether to use shortcut connections.
        g (int): Groups for convolutions.
        e (float): Expansion ratio.
    """
    super().__init__(c1, c2, n, shortcut, g, e)
    self.c_ = int(c2 * e)
    self.m = nn.Sequential(*(Bottleneck(self.c_, self.c_, shortcut, g, k=((1, 3), (3, 1)), e=1) for _ in range(n)))





ultralytics.nn.modules.block.RepC3

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

Bases: Module

Rep C3.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
n int

Number of RepConv blocks.

3
e float

Expansion ratio.

1.0
Source code in ultralytics/nn/modules/block.py
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
def __init__(self, c1, c2, n=3, e=1.0):
    """
    Initialize CSP Bottleneck with a single convolution.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        n (int): Number of RepConv blocks.
        e (float): Expansion ratio.
    """
    super().__init__()
    c_ = int(c2 * e)  # hidden channels
    self.cv1 = Conv(c1, c_, 1, 1)
    self.cv2 = Conv(c1, c_, 1, 1)
    self.m = nn.Sequential(*[RepConv(c_, c_) for _ in range(n)])
    self.cv3 = Conv(c_, c2, 1, 1) if c_ != c2 else nn.Identity()

forward

forward(x)

Forward pass of RepC3 module.

Source code in ultralytics/nn/modules/block.py
379
380
381
def forward(self, x):
    """Forward pass of RepC3 module."""
    return self.cv3(self.m(self.cv1(x)) + self.cv2(x))





ultralytics.nn.modules.block.C3TR

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

Bases: C3

C3 module with TransformerBlock().

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
n int

Number of Transformer blocks.

1
shortcut bool

Whether to use shortcut connections.

True
g int

Groups for convolutions.

1
e float

Expansion ratio.

0.5
Source code in ultralytics/nn/modules/block.py
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
    """
    Initialize C3 module with TransformerBlock.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        n (int): Number of Transformer blocks.
        shortcut (bool): Whether to use shortcut connections.
        g (int): Groups for convolutions.
        e (float): Expansion ratio.
    """
    super().__init__(c1, c2, n, shortcut, g, e)
    c_ = int(c2 * e)
    self.m = TransformerBlock(c_, c_, 4, n)





ultralytics.nn.modules.block.C3Ghost

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

Bases: C3

C3 module with GhostBottleneck().

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
n int

Number of Ghost bottleneck blocks.

1
shortcut bool

Whether to use shortcut connections.

True
g int

Groups for convolutions.

1
e float

Expansion ratio.

0.5
Source code in ultralytics/nn/modules/block.py
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
    """
    Initialize C3 module with GhostBottleneck.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        n (int): Number of Ghost bottleneck blocks.
        shortcut (bool): Whether to use shortcut connections.
        g (int): Groups for convolutions.
        e (float): Expansion ratio.
    """
    super().__init__(c1, c2, n, shortcut, g, e)
    c_ = int(c2 * e)  # hidden channels
    self.m = nn.Sequential(*(GhostBottleneck(c_, c_) for _ in range(n)))





ultralytics.nn.modules.block.GhostBottleneck

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

Bases: Module

Ghost Bottleneck https://github.com/huawei-noah/Efficient-AI-Backbones.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
k int

Kernel size.

3
s int

Stride.

1
Source code in ultralytics/nn/modules/block.py
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
def __init__(self, c1, c2, k=3, s=1):
    """
    Initialize Ghost Bottleneck module.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        k (int): Kernel size.
        s (int): Stride.
    """
    super().__init__()
    c_ = c2 // 2
    self.conv = nn.Sequential(
        GhostConv(c1, c_, 1, 1),  # pw
        DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(),  # dw
        GhostConv(c_, c2, 1, 1, act=False),  # pw-linear
    )
    self.shortcut = (
        nn.Sequential(DWConv(c1, c1, k, s, act=False), Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity()
    )

forward

forward(x)

Apply skip connection and concatenation to input tensor.

Source code in ultralytics/nn/modules/block.py
448
449
450
def forward(self, x):
    """Apply skip connection and concatenation to input tensor."""
    return self.conv(x) + self.shortcut(x)





ultralytics.nn.modules.block.Bottleneck

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

Bases: Module

Standard bottleneck.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
shortcut bool

Whether to use shortcut connection.

True
g int

Groups for convolutions.

1
k Tuple[int, int]

Kernel sizes for convolutions.

(3, 3)
e float

Expansion ratio.

0.5
Source code in ultralytics/nn/modules/block.py
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):
    """
    Initialize a standard bottleneck module.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        shortcut (bool): Whether to use shortcut connection.
        g (int): Groups for convolutions.
        k (Tuple[int, int]): Kernel sizes for convolutions.
        e (float): Expansion ratio.
    """
    super().__init__()
    c_ = int(c2 * e)  # hidden channels
    self.cv1 = Conv(c1, c_, k[0], 1)
    self.cv2 = Conv(c_, c2, k[1], 1, g=g)
    self.add = shortcut and c1 == c2

forward

forward(x)

Apply bottleneck with optional shortcut connection.

Source code in ultralytics/nn/modules/block.py
474
475
476
def forward(self, x):
    """Apply bottleneck with optional shortcut connection."""
    return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))





ultralytics.nn.modules.block.BottleneckCSP

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

Bases: Module

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

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
n int

Number of Bottleneck blocks.

1
shortcut bool

Whether to use shortcut connections.

True
g int

Groups for convolutions.

1
e float

Expansion ratio.

0.5
Source code in ultralytics/nn/modules/block.py
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
    """
    Initialize CSP Bottleneck.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        n (int): Number of Bottleneck blocks.
        shortcut (bool): Whether to use shortcut connections.
        g (int): Groups for convolutions.
        e (float): Expansion ratio.
    """
    super().__init__()
    c_ = int(c2 * e)  # hidden channels
    self.cv1 = Conv(c1, c_, 1, 1)
    self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
    self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
    self.cv4 = Conv(2 * c_, c2, 1, 1)
    self.bn = nn.BatchNorm2d(2 * c_)  # applied to cat(cv2, cv3)
    self.act = nn.SiLU()
    self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))

forward

forward(x)

Apply CSP bottleneck with 3 convolutions.

Source code in ultralytics/nn/modules/block.py
504
505
506
507
508
def forward(self, x):
    """Apply CSP bottleneck with 3 convolutions."""
    y1 = self.cv3(self.m(self.cv1(x)))
    y2 = self.cv2(x)
    return self.cv4(self.act(self.bn(torch.cat((y1, y2), 1))))





ultralytics.nn.modules.block.ResNetBlock

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

Bases: Module

ResNet block with standard convolution layers.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
s int

Stride.

1
e int

Expansion ratio.

4
Source code in ultralytics/nn/modules/block.py
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
def __init__(self, c1, c2, s=1, e=4):
    """
    Initialize ResNet block.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        s (int): Stride.
        e (int): Expansion ratio.
    """
    super().__init__()
    c3 = e * c2
    self.cv1 = Conv(c1, c2, k=1, s=1, act=True)
    self.cv2 = Conv(c2, c2, k=3, s=s, p=1, act=True)
    self.cv3 = Conv(c2, c3, k=1, act=False)
    self.shortcut = nn.Sequential(Conv(c1, c3, k=1, s=s, act=False)) if s != 1 or c1 != c3 else nn.Identity()

forward

forward(x)

Forward pass through the ResNet block.

Source code in ultralytics/nn/modules/block.py
531
532
533
def forward(self, x):
    """Forward pass through the ResNet block."""
    return F.relu(self.cv3(self.cv2(self.cv1(x))) + self.shortcut(x))





ultralytics.nn.modules.block.ResNetLayer

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

Bases: Module

ResNet layer with multiple ResNet blocks.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
s int

Stride.

1
is_first bool

Whether this is the first layer.

False
n int

Number of ResNet blocks.

1
e int

Expansion ratio.

4
Source code in ultralytics/nn/modules/block.py
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
def __init__(self, c1, c2, s=1, is_first=False, n=1, e=4):
    """
    Initialize ResNet layer.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        s (int): Stride.
        is_first (bool): Whether this is the first layer.
        n (int): Number of ResNet blocks.
        e (int): Expansion ratio.
    """
    super().__init__()
    self.is_first = is_first

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

forward

forward(x)

Forward pass through the ResNet layer.

Source code in ultralytics/nn/modules/block.py
563
564
565
def forward(self, x):
    """Forward pass through the ResNet layer."""
    return self.layer(x)





ultralytics.nn.modules.block.MaxSigmoidAttnBlock

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

Bases: Module

Max Sigmoid attention block.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
nh int

Number of heads.

1
ec int

Embedding channels.

128
gc int

Guide channels.

512
scale bool

Whether to use learnable scale parameter.

False
Source code in ultralytics/nn/modules/block.py
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
def __init__(self, c1, c2, nh=1, ec=128, gc=512, scale=False):
    """
    Initialize MaxSigmoidAttnBlock.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        nh (int): Number of heads.
        ec (int): Embedding channels.
        gc (int): Guide channels.
        scale (bool): Whether to use learnable scale parameter.
    """
    super().__init__()
    self.nh = nh
    self.hc = c2 // nh
    self.ec = Conv(c1, ec, k=1, act=False) if c1 != ec else None
    self.gl = nn.Linear(gc, ec)
    self.bias = nn.Parameter(torch.zeros(nh))
    self.proj_conv = Conv(c1, c2, k=3, s=1, act=False)
    self.scale = nn.Parameter(torch.ones(1, nh, 1, 1)) if scale else 1.0

forward

forward(x, guide)

Forward pass of MaxSigmoidAttnBlock.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

required
guide Tensor

Guide tensor.

required

Returns:

Type Description
Tensor

Output tensor after attention.

Source code in ultralytics/nn/modules/block.py
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
def forward(self, x, guide):
    """
    Forward pass of MaxSigmoidAttnBlock.

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

    Returns:
        (torch.Tensor): Output tensor after attention.
    """
    bs, _, h, w = x.shape

    guide = self.gl(guide)
    guide = guide.view(bs, guide.shape[1], self.nh, self.hc)
    embed = self.ec(x) if self.ec is not None else x
    embed = embed.view(bs, self.nh, self.hc, h, w)

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

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





ultralytics.nn.modules.block.C2fAttn

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

Bases: Module

C2f module with an additional attn module.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
n int

Number of Bottleneck blocks.

1
ec int

Embedding channels for attention.

128
nh int

Number of heads for attention.

1
gc int

Guide channels for attention.

512
shortcut bool

Whether to use shortcut connections.

False
g int

Groups for convolutions.

1
e float

Expansion ratio.

0.5
Source code in ultralytics/nn/modules/block.py
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
def __init__(self, c1, c2, n=1, ec=128, nh=1, gc=512, shortcut=False, g=1, e=0.5):
    """
    Initialize C2f module with attention mechanism.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        n (int): Number of Bottleneck blocks.
        ec (int): Embedding channels for attention.
        nh (int): Number of heads for attention.
        gc (int): Guide channels for attention.
        shortcut (bool): Whether to use shortcut connections.
        g (int): Groups for convolutions.
        e (float): Expansion ratio.
    """
    super().__init__()
    self.c = int(c2 * e)  # hidden channels
    self.cv1 = Conv(c1, 2 * self.c, 1, 1)
    self.cv2 = Conv((3 + n) * self.c, c2, 1)  # optional act=FReLU(c2)
    self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))
    self.attn = MaxSigmoidAttnBlock(self.c, self.c, gc=gc, ec=ec, nh=nh)

forward

forward(x, guide)

Forward pass through C2f layer with attention.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

required
guide Tensor

Guide tensor for attention.

required

Returns:

Type Description
Tensor

Output tensor after processing.

Source code in ultralytics/nn/modules/block.py
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
def forward(self, x, guide):
    """
    Forward pass through C2f layer with attention.

    Args:
        x (torch.Tensor): Input tensor.
        guide (torch.Tensor): Guide tensor for attention.

    Returns:
        (torch.Tensor): Output tensor after processing.
    """
    y = list(self.cv1(x).chunk(2, 1))
    y.extend(m(y[-1]) for m in self.m)
    y.append(self.attn(y[-1], guide))
    return self.cv2(torch.cat(y, 1))

forward_split

forward_split(x, guide)

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

Parameters:

Name Type Description Default
x Tensor

Input tensor.

required
guide Tensor

Guide tensor for attention.

required

Returns:

Type Description
Tensor

Output tensor after processing.

Source code in ultralytics/nn/modules/block.py
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
def forward_split(self, x, guide):
    """
    Forward pass using split() instead of chunk().

    Args:
        x (torch.Tensor): Input tensor.
        guide (torch.Tensor): Guide tensor for attention.

    Returns:
        (torch.Tensor): Output tensor after processing.
    """
    y = list(self.cv1(x).split((self.c, self.c), 1))
    y.extend(m(y[-1]) for m in self.m)
    y.append(self.attn(y[-1], guide))
    return self.cv2(torch.cat(y, 1))





ultralytics.nn.modules.block.ImagePoolingAttn

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

Bases: Module

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

Parameters:

Name Type Description Default
ec int

Embedding channels.

256
ch tuple

Channel dimensions for feature maps.

()
ct int

Channel dimension for text embeddings.

512
nh int

Number of attention heads.

8
k int

Kernel size for pooling.

3
scale bool

Whether to use learnable scale parameter.

False
Source code in ultralytics/nn/modules/block.py
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
def __init__(self, ec=256, ch=(), ct=512, nh=8, k=3, scale=False):
    """
    Initialize ImagePoolingAttn module.

    Args:
        ec (int): Embedding channels.
        ch (tuple): Channel dimensions for feature maps.
        ct (int): Channel dimension for text embeddings.
        nh (int): Number of attention heads.
        k (int): Kernel size for pooling.
        scale (bool): Whether to use learnable scale parameter.
    """
    super().__init__()

    nf = len(ch)
    self.query = nn.Sequential(nn.LayerNorm(ct), nn.Linear(ct, ec))
    self.key = nn.Sequential(nn.LayerNorm(ec), nn.Linear(ec, ec))
    self.value = nn.Sequential(nn.LayerNorm(ec), nn.Linear(ec, ec))
    self.proj = nn.Linear(ec, ct)
    self.scale = nn.Parameter(torch.tensor([0.0]), requires_grad=True) if scale else 1.0
    self.projections = nn.ModuleList([nn.Conv2d(in_channels, ec, kernel_size=1) for in_channels in ch])
    self.im_pools = nn.ModuleList([nn.AdaptiveMaxPool2d((k, k)) for _ in range(nf)])
    self.ec = ec
    self.nh = nh
    self.nf = nf
    self.hc = ec // nh
    self.k = k

forward

forward(x, text)

Forward pass of ImagePoolingAttn.

Parameters:

Name Type Description Default
x List[Tensor]

List of input feature maps.

required
text Tensor

Text embeddings.

required

Returns:

Type Description
Tensor

Enhanced text embeddings.

Source code in ultralytics/nn/modules/block.py
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
def forward(self, x, text):
    """
    Forward pass of ImagePoolingAttn.

    Args:
        x (List[torch.Tensor]): List of input feature maps.
        text (torch.Tensor): Text embeddings.

    Returns:
        (torch.Tensor): Enhanced text embeddings.
    """
    bs = x[0].shape[0]
    assert len(x) == self.nf
    num_patches = self.k**2
    x = [pool(proj(x)).view(bs, -1, num_patches) for (x, proj, pool) in zip(x, self.projections, self.im_pools)]
    x = torch.cat(x, dim=-1).transpose(1, 2)
    q = self.query(text)
    k = self.key(x)
    v = self.value(x)

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

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

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





ultralytics.nn.modules.block.ContrastiveHead

ContrastiveHead()

Bases: Module

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

Source code in ultralytics/nn/modules/block.py
748
749
750
751
752
753
def __init__(self):
    """Initialize ContrastiveHead with region-text similarity parameters."""
    super().__init__()
    # NOTE: use -10.0 to keep the init cls loss consistency with other losses
    self.bias = nn.Parameter(torch.tensor([-10.0]))
    self.logit_scale = nn.Parameter(torch.ones([]) * torch.tensor(1 / 0.07).log())

forward

forward(x, w)

Forward function of contrastive learning.

Parameters:

Name Type Description Default
x Tensor

Image features.

required
w Tensor

Text features.

required

Returns:

Type Description
Tensor

Similarity scores.

Source code in ultralytics/nn/modules/block.py
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
def forward(self, x, w):
    """
    Forward function of contrastive learning.

    Args:
        x (torch.Tensor): Image features.
        w (torch.Tensor): Text features.

    Returns:
        (torch.Tensor): Similarity scores.
    """
    x = F.normalize(x, dim=1, p=2)
    w = F.normalize(w, dim=-1, p=2)
    x = torch.einsum("bchw,bkc->bkhw", x, w)
    return x * self.logit_scale.exp() + self.bias





ultralytics.nn.modules.block.BNContrastiveHead

BNContrastiveHead(embed_dims: int)

Bases: Module

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

Parameters:

Name Type Description Default
embed_dims int

Embed dimensions of text and image features.

required

Parameters:

Name Type Description Default
embed_dims int

Embedding dimensions for features.

required
Source code in ultralytics/nn/modules/block.py
780
781
782
783
784
785
786
787
788
789
790
791
792
def __init__(self, embed_dims: int):
    """
    Initialize BNContrastiveHead.

    Args:
        embed_dims (int): Embedding dimensions for features.
    """
    super().__init__()
    self.norm = nn.BatchNorm2d(embed_dims)
    # NOTE: use -10.0 to keep the init cls loss consistency with other losses
    self.bias = nn.Parameter(torch.tensor([-10.0]))
    # use -1.0 is more stable
    self.logit_scale = nn.Parameter(-1.0 * torch.ones([]))

forward

forward(x, w)

Forward function of contrastive learning with batch normalization.

Parameters:

Name Type Description Default
x Tensor

Image features.

required
w Tensor

Text features.

required

Returns:

Type Description
Tensor

Similarity scores.

Source code in ultralytics/nn/modules/block.py
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
def forward(self, x, w):
    """
    Forward function of contrastive learning with batch normalization.

    Args:
        x (torch.Tensor): Image features.
        w (torch.Tensor): Text features.

    Returns:
        (torch.Tensor): Similarity scores.
    """
    x = self.norm(x)
    w = F.normalize(w, dim=-1, p=2)

    x = torch.einsum("bchw,bkc->bkhw", x, w)
    return x * self.logit_scale.exp() + self.bias

forward_fuse

forward_fuse(x, w)

Passes input out unchanged.

TODO: Update or remove?

Source code in ultralytics/nn/modules/block.py
801
802
803
804
805
806
807
def forward_fuse(self, x, w):
    """
    Passes input out unchanged.

    TODO: Update or remove?
    """
    return x

fuse

fuse()

Fuse the batch normalization layer in the BNContrastiveHead module.

Source code in ultralytics/nn/modules/block.py
794
795
796
797
798
799
def fuse(self):
    """Fuse the batch normalization layer in the BNContrastiveHead module."""
    del self.norm
    del self.bias
    del self.logit_scale
    self.forward = self.forward_fuse





ultralytics.nn.modules.block.RepBottleneck

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

Bases: Bottleneck

Rep bottleneck.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
shortcut bool

Whether to use shortcut connection.

True
g int

Groups for convolutions.

1
k Tuple[int, int]

Kernel sizes for convolutions.

(3, 3)
e float

Expansion ratio.

0.5
Source code in ultralytics/nn/modules/block.py
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):
    """
    Initialize RepBottleneck.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        shortcut (bool): Whether to use shortcut connection.
        g (int): Groups for convolutions.
        k (Tuple[int, int]): Kernel sizes for convolutions.
        e (float): Expansion ratio.
    """
    super().__init__(c1, c2, shortcut, g, k, e)
    c_ = int(c2 * e)  # hidden channels
    self.cv1 = RepConv(c1, c_, k[0], 1)





ultralytics.nn.modules.block.RepCSP

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

Bases: C3

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

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
n int

Number of RepBottleneck blocks.

1
shortcut bool

Whether to use shortcut connections.

True
g int

Groups for convolutions.

1
e float

Expansion ratio.

0.5
Source code in ultralytics/nn/modules/block.py
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
    """
    Initialize RepCSP layer.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        n (int): Number of RepBottleneck blocks.
        shortcut (bool): Whether to use shortcut connections.
        g (int): Groups for convolutions.
        e (float): Expansion ratio.
    """
    super().__init__(c1, c2, n, shortcut, g, e)
    c_ = int(c2 * e)  # hidden channels
    self.m = nn.Sequential(*(RepBottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))





ultralytics.nn.modules.block.RepNCSPELAN4

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

Bases: Module

CSP-ELAN.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
c3 int

Intermediate channels.

required
c4 int

Intermediate channels for RepCSP.

required
n int

Number of RepCSP blocks.

1
Source code in ultralytics/nn/modules/block.py
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
def __init__(self, c1, c2, c3, c4, n=1):
    """
    Initialize CSP-ELAN layer.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        c3 (int): Intermediate channels.
        c4 (int): Intermediate channels for RepCSP.
        n (int): Number of RepCSP blocks.
    """
    super().__init__()
    self.c = c3 // 2
    self.cv1 = Conv(c1, c3, 1, 1)
    self.cv2 = nn.Sequential(RepCSP(c3 // 2, c4, n), Conv(c4, c4, 3, 1))
    self.cv3 = nn.Sequential(RepCSP(c4, c4, n), Conv(c4, c4, 3, 1))
    self.cv4 = Conv(c3 + (2 * c4), c2, 1, 1)

forward

forward(x)

Forward pass through RepNCSPELAN4 layer.

Source code in ultralytics/nn/modules/block.py
888
889
890
891
892
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
894
895
896
897
898
def forward_split(self, x):
    """Forward pass using split() instead of chunk()."""
    y = list(self.cv1(x).split((self.c, self.c), 1))
    y.extend(m(y[-1]) for m in [self.cv2, self.cv3])
    return self.cv4(torch.cat(y, 1))





ultralytics.nn.modules.block.ELAN1

ELAN1(c1, c2, c3, c4)

Bases: RepNCSPELAN4

ELAN1 module with 4 convolutions.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
c3 int

Intermediate channels.

required
c4 int

Intermediate channels for convolutions.

required
Source code in ultralytics/nn/modules/block.py
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
def __init__(self, c1, c2, c3, c4):
    """
    Initialize ELAN1 layer.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        c3 (int): Intermediate channels.
        c4 (int): Intermediate channels for convolutions.
    """
    super().__init__(c1, c2, c3, c4)
    self.c = c3 // 2
    self.cv1 = Conv(c1, c3, 1, 1)
    self.cv2 = Conv(c3 // 2, c4, 3, 1)
    self.cv3 = Conv(c4, c4, 3, 1)
    self.cv4 = Conv(c3 + (2 * c4), c2, 1, 1)





ultralytics.nn.modules.block.AConv

AConv(c1, c2)

Bases: Module

AConv.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
Source code in ultralytics/nn/modules/block.py
925
926
927
928
929
930
931
932
933
934
def __init__(self, c1, c2):
    """
    Initialize AConv module.

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

forward

forward(x)

Forward pass through AConv layer.

Source code in ultralytics/nn/modules/block.py
936
937
938
939
def forward(self, x):
    """Forward pass through AConv layer."""
    x = torch.nn.functional.avg_pool2d(x, 2, 1, 0, False, True)
    return self.cv1(x)





ultralytics.nn.modules.block.ADown

ADown(c1, c2)

Bases: Module

ADown.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
Source code in ultralytics/nn/modules/block.py
945
946
947
948
949
950
951
952
953
954
955
956
def __init__(self, c1, c2):
    """
    Initialize ADown module.

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

forward

forward(x)

Forward pass through ADown layer.

Source code in ultralytics/nn/modules/block.py
958
959
960
961
962
963
964
965
def forward(self, x):
    """Forward pass through ADown layer."""
    x = torch.nn.functional.avg_pool2d(x, 2, 1, 0, False, True)
    x1, x2 = x.chunk(2, 1)
    x1 = self.cv1(x1)
    x2 = torch.nn.functional.max_pool2d(x2, 3, 2, 1)
    x2 = self.cv2(x2)
    return torch.cat((x1, x2), 1)





ultralytics.nn.modules.block.SPPELAN

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

Bases: Module

SPP-ELAN.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
c3 int

Intermediate channels.

required
k int

Kernel size for max pooling.

5
Source code in ultralytics/nn/modules/block.py
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
def __init__(self, c1, c2, c3, k=5):
    """
    Initialize SPP-ELAN block.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        c3 (int): Intermediate channels.
        k (int): Kernel size for max pooling.
    """
    super().__init__()
    self.c = c3
    self.cv1 = Conv(c1, c3, 1, 1)
    self.cv2 = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
    self.cv3 = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
    self.cv4 = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
    self.cv5 = Conv(4 * c3, c2, 1, 1)

forward

forward(x)

Forward pass through SPPELAN layer.

Source code in ultralytics/nn/modules/block.py
989
990
991
992
993
def forward(self, x):
    """Forward pass through SPPELAN layer."""
    y = [self.cv1(x)]
    y.extend(m(y[-1]) for m in [self.cv2, self.cv3, self.cv4])
    return self.cv5(torch.cat(y, 1))





ultralytics.nn.modules.block.CBLinear

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

Bases: Module

CBLinear.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2s List[int]

List of output channel sizes.

required
k int

Kernel size.

1
s int

Stride.

1
p int | None

Padding.

None
g int

Groups.

1
Source code in ultralytics/nn/modules/block.py
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
def __init__(self, c1, c2s, k=1, s=1, p=None, g=1):
    """
    Initialize CBLinear module.

    Args:
        c1 (int): Input channels.
        c2s (List[int]): List of output channel sizes.
        k (int): Kernel size.
        s (int): Stride.
        p (int | None): Padding.
        g (int): Groups.
    """
    super().__init__()
    self.c2s = c2s
    self.conv = nn.Conv2d(c1, sum(c2s), k, s, autopad(k, p), groups=g, bias=True)

forward

forward(x)

Forward pass through CBLinear layer.

Source code in ultralytics/nn/modules/block.py
1015
1016
1017
def forward(self, x):
    """Forward pass through CBLinear layer."""
    return self.conv(x).split(self.c2s, dim=1)





ultralytics.nn.modules.block.CBFuse

CBFuse(idx)

Bases: Module

CBFuse.

Parameters:

Name Type Description Default
idx List[int]

Indices for feature selection.

required
Source code in ultralytics/nn/modules/block.py
1023
1024
1025
1026
1027
1028
1029
1030
1031
def __init__(self, idx):
    """
    Initialize CBFuse module.

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

forward

forward(xs)

Forward pass through CBFuse layer.

Parameters:

Name Type Description Default
xs List[Tensor]

List of input tensors.

required

Returns:

Type Description
Tensor

Fused output tensor.

Source code in ultralytics/nn/modules/block.py
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
def forward(self, xs):
    """
    Forward pass through CBFuse layer.

    Args:
        xs (List[torch.Tensor]): List of input tensors.

    Returns:
        (torch.Tensor): Fused output tensor.
    """
    target_size = xs[-1].shape[2:]
    res = [F.interpolate(x[self.idx[i]], size=target_size, mode="nearest") for i, x in enumerate(xs[:-1])]
    return torch.sum(torch.stack(res + xs[-1:]), dim=0)





ultralytics.nn.modules.block.C3f

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

Bases: Module

Faster Implementation of CSP Bottleneck with 2 convolutions.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
n int

Number of Bottleneck blocks.

1
shortcut bool

Whether to use shortcut connections.

False
g int

Groups for convolutions.

1
e float

Expansion ratio.

0.5
Source code in ultralytics/nn/modules/block.py
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
    """
    Initialize CSP bottleneck layer with two convolutions.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        n (int): Number of Bottleneck blocks.
        shortcut (bool): Whether to use shortcut connections.
        g (int): Groups for convolutions.
        e (float): Expansion ratio.
    """
    super().__init__()
    c_ = int(c2 * e)  # hidden channels
    self.cv1 = Conv(c1, c_, 1, 1)
    self.cv2 = Conv(c1, c_, 1, 1)
    self.cv3 = Conv((2 + n) * c_, c2, 1)  # optional act=FReLU(c2)
    self.m = nn.ModuleList(Bottleneck(c_, c_, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))

forward

forward(x)

Forward pass through C3f layer.

Source code in ultralytics/nn/modules/block.py
1070
1071
1072
1073
1074
def forward(self, x):
    """Forward pass through C3f layer."""
    y = [self.cv2(x), self.cv1(x)]
    y.extend(m(y[-1]) for m in self.m)
    return self.cv3(torch.cat(y, 1))





ultralytics.nn.modules.block.C3k2

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

Bases: C2f

Faster Implementation of CSP Bottleneck with 2 convolutions.

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
n int

Number of blocks.

1
c3k bool

Whether to use C3k blocks.

False
e float

Expansion ratio.

0.5
g int

Groups for convolutions.

1
shortcut bool

Whether to use shortcut connections.

True
Source code in ultralytics/nn/modules/block.py
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
def __init__(self, c1, c2, n=1, c3k=False, e=0.5, g=1, shortcut=True):
    """
    Initialize C3k2 module.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        n (int): Number of blocks.
        c3k (bool): Whether to use C3k blocks.
        e (float): Expansion ratio.
        g (int): Groups for convolutions.
        shortcut (bool): Whether to use shortcut connections.
    """
    super().__init__(c1, c2, n, shortcut, g, e)
    self.m = nn.ModuleList(
        C3k(self.c, self.c, 2, shortcut, g) if c3k else Bottleneck(self.c, self.c, shortcut, g) for _ in range(n)
    )





ultralytics.nn.modules.block.C3k

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

Bases: C3

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

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
n int

Number of Bottleneck blocks.

1
shortcut bool

Whether to use shortcut connections.

True
g int

Groups for convolutions.

1
e float

Expansion ratio.

0.5
k int

Kernel size.

3
Source code in ultralytics/nn/modules/block.py
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, k=3):
    """
    Initialize C3k module.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        n (int): Number of Bottleneck blocks.
        shortcut (bool): Whether to use shortcut connections.
        g (int): Groups for convolutions.
        e (float): Expansion ratio.
        k (int): Kernel size.
    """
    super().__init__(c1, c2, n, shortcut, g, e)
    c_ = int(c2 * e)  # hidden channels
    # self.m = nn.Sequential(*(RepBottleneck(c_, c_, shortcut, g, k=(k, k), e=1.0) for _ in range(n)))
    self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, k=(k, k), e=1.0) for _ in range(n)))





ultralytics.nn.modules.block.RepVGGDW

RepVGGDW(ed)

Bases: Module

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

Parameters:

Name Type Description Default
ed int

Input and output channels.

required
Source code in ultralytics/nn/modules/block.py
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
def __init__(self, ed) -> None:
    """
    Initialize RepVGGDW module.

    Args:
        ed (int): Input and output channels.
    """
    super().__init__()
    self.conv = Conv(ed, ed, 7, 1, 3, g=ed, act=False)
    self.conv1 = Conv(ed, ed, 3, 1, 1, g=ed, act=False)
    self.dim = ed
    self.act = nn.SiLU()

forward

forward(x)

Perform a forward pass of the RepVGGDW block.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

required

Returns:

Type Description
Tensor

Output tensor after applying the depth wise separable convolution.

Source code in ultralytics/nn/modules/block.py
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
def forward(self, x):
    """
    Perform a forward pass of the RepVGGDW block.

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

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

forward_fuse

forward_fuse(x)

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

Parameters:

Name Type Description Default
x Tensor

Input tensor.

required

Returns:

Type Description
Tensor

Output tensor after applying the depth wise separable convolution.

Source code in ultralytics/nn/modules/block.py
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
def forward_fuse(self, x):
    """
    Perform a forward pass of the RepVGGDW block without fusing the convolutions.

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

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

fuse

fuse()

Fuse the convolutional layers in the RepVGGDW block.

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

Source code in ultralytics/nn/modules/block.py
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
@torch.no_grad()
def fuse(self):
    """
    Fuse the convolutional layers in the RepVGGDW block.

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

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

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

    final_conv_w = conv_w + conv1_w
    final_conv_b = conv_b + conv1_b

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

    self.conv = conv
    del self.conv1





ultralytics.nn.modules.block.CIB

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

Bases: Module

Conditional Identity Block (CIB) module.

Parameters:

Name Type Description Default
c1 int

Number of input channels.

required
c2 int

Number of output channels.

required
shortcut bool

Whether to add a shortcut connection. Defaults to True.

True
e float

Scaling factor for the hidden channels. Defaults to 0.5.

0.5
lk bool

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

False

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
shortcut bool

Whether to use shortcut connection.

True
e float

Expansion ratio.

0.5
lk bool

Whether to use RepVGGDW.

False
Source code in ultralytics/nn/modules/block.py
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
def __init__(self, c1, c2, shortcut=True, e=0.5, lk=False):
    """
    Initialize the CIB module.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        shortcut (bool): Whether to use shortcut connection.
        e (float): Expansion ratio.
        lk (bool): Whether to use RepVGGDW.
    """
    super().__init__()
    c_ = int(c2 * e)  # hidden channels
    self.cv1 = nn.Sequential(
        Conv(c1, c1, 3, g=c1),
        Conv(c1, 2 * c_, 1),
        RepVGGDW(2 * c_) if lk else Conv(2 * c_, 2 * c_, 3, g=2 * c_),
        Conv(2 * c_, c2, 1),
        Conv(c2, c2, 3, g=c2),
    )

    self.add = shortcut and c1 == c2

forward

forward(x)

Forward pass of the CIB module.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

required

Returns:

Type Description
Tensor

Output tensor.

Source code in ultralytics/nn/modules/block.py
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
def forward(self, x):
    """
    Forward pass of the CIB module.

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

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





ultralytics.nn.modules.block.C2fCIB

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

Bases: C2f

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

Parameters:

Name Type Description Default
c1 int

Number of input channels.

required
c2 int

Number of output channels.

required
n int

Number of CIB modules to stack. Defaults to 1.

1
shortcut bool

Whether to use shortcut connection. Defaults to False.

False
lk bool

Whether to use local key connection. Defaults to False.

False
g int

Number of groups for grouped convolution. Defaults to 1.

1
e float

Expansion ratio for CIB modules. Defaults to 0.5.

0.5

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
n int

Number of CIB modules.

1
shortcut bool

Whether to use shortcut connection.

False
lk bool

Whether to use local key connection.

False
g int

Groups for convolutions.

1
e float

Expansion ratio.

0.5
Source code in ultralytics/nn/modules/block.py
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
def __init__(self, c1, c2, n=1, shortcut=False, lk=False, g=1, e=0.5):
    """
    Initialize C2fCIB module.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        n (int): Number of CIB modules.
        shortcut (bool): Whether to use shortcut connection.
        lk (bool): Whether to use local key connection.
        g (int): Groups for convolutions.
        e (float): Expansion ratio.
    """
    super().__init__(c1, c2, n, shortcut, g, e)
    self.m = nn.ModuleList(CIB(self.c, self.c, shortcut, e=1.0, lk=lk) for _ in range(n))





ultralytics.nn.modules.block.Attention

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

Bases: Module

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

Parameters:

Name Type Description Default
dim int

The input tensor dimension.

required
num_heads int

The number of attention heads.

8
attn_ratio float

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

0.5

Attributes:

Name Type Description
num_heads int

The number of attention heads.

head_dim int

The dimension of each attention head.

key_dim int

The dimension of the attention key.

scale float

The scaling factor for the attention scores.

qkv Conv

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

proj Conv

Convolutional layer for projecting the attended values.

pe Conv

Convolutional layer for positional encoding.

Parameters:

Name Type Description Default
dim int

Input dimension.

required
num_heads int

Number of attention heads.

8
attn_ratio float

Attention ratio for key dimension.

0.5
Source code in ultralytics/nn/modules/block.py
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
def __init__(self, dim, num_heads=8, attn_ratio=0.5):
    """
    Initialize multi-head attention module.

    Args:
        dim (int): Input dimension.
        num_heads (int): Number of attention heads.
        attn_ratio (float): Attention ratio for key dimension.
    """
    super().__init__()
    self.num_heads = num_heads
    self.head_dim = dim // num_heads
    self.key_dim = int(self.head_dim * attn_ratio)
    self.scale = self.key_dim**-0.5
    nh_kd = self.key_dim * num_heads
    h = dim + nh_kd * 2
    self.qkv = Conv(dim, h, 1, act=False)
    self.proj = Conv(dim, dim, 1, act=False)
    self.pe = Conv(dim, dim, 3, 1, g=dim, act=False)

forward

forward(x)

Forward pass of the Attention module.

Parameters:

Name Type Description Default
x Tensor

The input tensor.

required

Returns:

Type Description
Tensor

The output tensor after self-attention.

Source code in ultralytics/nn/modules/block.py
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
def forward(self, x):
    """
    Forward pass of the Attention module.

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

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

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





ultralytics.nn.modules.block.PSABlock

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

Bases: Module

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

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

Attributes:

Name Type Description
attn Attention

Multi-head attention module.

ffn Sequential

Feed-forward neural network module.

add bool

Flag indicating whether to add shortcut connections.

Methods:

Name Description
forward

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

Examples:

Create a PSABlock and perform a forward pass

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

Parameters:

Name Type Description Default
c int

Input and output channels.

required
attn_ratio float

Attention ratio for key dimension.

0.5
num_heads int

Number of attention heads.

4
shortcut bool

Whether to use shortcut connections.

True
Source code in ultralytics/nn/modules/block.py
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
def __init__(self, c, attn_ratio=0.5, num_heads=4, shortcut=True) -> None:
    """
    Initialize the PSABlock.

    Args:
        c (int): Input and output channels.
        attn_ratio (float): Attention ratio for key dimension.
        num_heads (int): Number of attention heads.
        shortcut (bool): Whether to use shortcut connections.
    """
    super().__init__()

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

forward

forward(x)

Execute a forward pass through PSABlock.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

required

Returns:

Type Description
Tensor

Output tensor after attention and feed-forward processing.

Source code in ultralytics/nn/modules/block.py
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
def forward(self, x):
    """
    Execute a forward pass through PSABlock.

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

    Returns:
        (torch.Tensor): Output tensor after attention and feed-forward processing.
    """
    x = x + self.attn(x) if self.add else self.attn(x)
    x = x + self.ffn(x) if self.add else self.ffn(x)
    return x





ultralytics.nn.modules.block.PSA

PSA(c1, c2, e=0.5)

Bases: Module

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

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

Attributes:

Name Type Description
c int

Number of hidden channels after applying the initial convolution.

cv1 Conv

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

cv2 Conv

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

attn Attention

Attention module for position-sensitive attention.

ffn Sequential

Feed-forward network for further processing.

Methods:

Name Description
forward

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

Examples:

Create a PSA module and apply it to an input tensor

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

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
e float

Expansion ratio.

0.5
Source code in ultralytics/nn/modules/block.py
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
def __init__(self, c1, c2, e=0.5):
    """
    Initialize PSA module.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        e (float): Expansion ratio.
    """
    super().__init__()
    assert c1 == c2
    self.c = int(c1 * e)
    self.cv1 = Conv(c1, 2 * self.c, 1, 1)
    self.cv2 = Conv(2 * self.c, c1, 1)

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

forward

forward(x)

Execute forward pass in PSA module.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

required

Returns:

Type Description
Tensor

Output tensor after attention and feed-forward processing.

Source code in ultralytics/nn/modules/block.py
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
def forward(self, x):
    """
    Execute forward pass in PSA module.

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

    Returns:
        (torch.Tensor): Output tensor after attention and feed-forward processing.
    """
    a, b = self.cv1(x).split((self.c, self.c), dim=1)
    b = b + self.attn(b)
    b = b + self.ffn(b)
    return self.cv2(torch.cat((a, b), 1))





ultralytics.nn.modules.block.C2PSA

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

Bases: Module

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

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

Attributes:

Name Type Description
c int

Number of hidden channels.

cv1 Conv

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

cv2 Conv

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

m Sequential

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

Methods:

Name Description
forward

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

Notes

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

Examples:

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

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
n int

Number of PSABlock modules.

1
e float

Expansion ratio.

0.5
Source code in ultralytics/nn/modules/block.py
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
def __init__(self, c1, c2, n=1, e=0.5):
    """
    Initialize C2PSA module.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        n (int): Number of PSABlock modules.
        e (float): Expansion ratio.
    """
    super().__init__()
    assert c1 == c2
    self.c = int(c1 * e)
    self.cv1 = Conv(c1, 2 * self.c, 1, 1)
    self.cv2 = Conv(2 * self.c, c1, 1)

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

forward

forward(x)

Process the input tensor through a series of PSA blocks.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

required

Returns:

Type Description
Tensor

Output tensor after processing.

Source code in ultralytics/nn/modules/block.py
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
def forward(self, x):
    """
    Process the input tensor through a series of PSA blocks.

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

    Returns:
        (torch.Tensor): Output tensor after processing.
    """
    a, b = self.cv1(x).split((self.c, self.c), dim=1)
    b = self.m(b)
    return self.cv2(torch.cat((a, b), 1))





ultralytics.nn.modules.block.C2fPSA

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

Bases: C2f

C2fPSA module with enhanced feature extraction using PSA blocks.

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

Attributes:

Name Type Description
c int

Number of hidden channels.

cv1 Conv

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

cv2 Conv

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

m ModuleList

List of PSA blocks for feature extraction.

Methods:

Name Description
forward

Performs a forward pass through the C2fPSA module.

forward_split

Performs a forward pass using split() instead of chunk().

Examples:

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

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
n int

Number of PSABlock modules.

1
e float

Expansion ratio.

0.5
Source code in ultralytics/nn/modules/block.py
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
def __init__(self, c1, c2, n=1, e=0.5):
    """
    Initialize C2fPSA module.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        n (int): Number of PSABlock modules.
        e (float): Expansion ratio.
    """
    assert c1 == c2
    super().__init__(c1, c2, n=n, e=e)
    self.m = nn.ModuleList(PSABlock(self.c, attn_ratio=0.5, num_heads=self.c // 64) for _ in range(n))





ultralytics.nn.modules.block.SCDown

SCDown(c1, c2, k, s)

Bases: Module

SCDown module for downsampling with separable convolutions.

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

Attributes:

Name Type Description
cv1 Conv

Pointwise convolution layer that reduces the number of channels.

cv2 Conv

Depthwise convolution layer that performs spatial downsampling.

Methods:

Name Description
forward

Applies the SCDown module to the input tensor.

Examples:

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

Parameters:

Name Type Description Default
c1 int

Input channels.

required
c2 int

Output channels.

required
k int

Kernel size.

required
s int

Stride.

required
Source code in ultralytics/nn/modules/block.py
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
def __init__(self, c1, c2, k, s):
    """
    Initialize SCDown module.

    Args:
        c1 (int): Input channels.
        c2 (int): Output channels.
        k (int): Kernel size.
        s (int): Stride.
    """
    super().__init__()
    self.cv1 = Conv(c1, c2, 1, 1)
    self.cv2 = Conv(c2, c2, k=k, s=s, g=c2, act=False)

forward

forward(x)

Apply convolution and downsampling to the input tensor.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

required

Returns:

Type Description
Tensor

Downsampled output tensor.

Source code in ultralytics/nn/modules/block.py
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
def forward(self, x):
    """
    Apply convolution and downsampling to the input tensor.

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

    Returns:
        (torch.Tensor): Downsampled output tensor.
    """
    return self.cv2(self.cv1(x))





ultralytics.nn.modules.block.TorchVision

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

Bases: Module

TorchVision module to allow loading any torchvision model.

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

Attributes:

Name Type Description
m Module

The loaded torchvision model, possibly truncated and unwrapped.

Parameters:

Name Type Description Default
model str

Name of the torchvision model to load.

required
weights str

Pre-trained weights to load. Default is "DEFAULT".

'DEFAULT'
unwrap bool

If True, unwraps the model to a sequential containing all but the last truncate layers. Default is True.

True
truncate int

Number of layers to truncate from the end if unwrap is True. Default is 2.

2
split bool

Returns output from intermediate child modules as list. Default is False.

False

Parameters:

Name Type Description Default
model str

Name of the torchvision model to load.

required
weights str

Pre-trained weights to load.

'DEFAULT'
unwrap bool

Whether to unwrap the model.

True
truncate int

Number of layers to truncate.

2
split bool

Whether to split the output.

False
Source code in ultralytics/nn/modules/block.py
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
def __init__(self, model, weights="DEFAULT", unwrap=True, truncate=2, split=False):
    """
    Load the model and weights from torchvision.

    Args:
        model (str): Name of the torchvision model to load.
        weights (str): Pre-trained weights to load.
        unwrap (bool): Whether to unwrap the model.
        truncate (int): Number of layers to truncate.
        split (bool): Whether to split the output.
    """
    import torchvision  # scope for faster 'import ultralytics'

    super().__init__()
    if hasattr(torchvision.models, "get_model"):
        self.m = torchvision.models.get_model(model, weights=weights)
    else:
        self.m = torchvision.models.__dict__[model](pretrained=bool(weights))
    if unwrap:
        layers = list(self.m.children())
        if isinstance(layers[0], nn.Sequential):  # Second-level for some models like EfficientNet, Swin
            layers = [*list(layers[0].children()), *layers[1:]]
        self.m = nn.Sequential(*(layers[:-truncate] if truncate else layers))
        self.split = split
    else:
        self.split = False
        self.m.head = self.m.heads = nn.Identity()

forward

forward(x)

Forward pass through the model.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

required

Returns:

Type Description
Tensor | List[Tensor]

Output tensor or list of tensors.

Source code in ultralytics/nn/modules/block.py
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
def forward(self, x):
    """
    Forward pass through the model.

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

    Returns:
        (torch.Tensor | List[torch.Tensor]): Output tensor or list of tensors.
    """
    if self.split:
        y = [x]
        y.extend(m(y[-1]) for m in self.m)
    else:
        y = self.m(x)
    return y





ultralytics.nn.modules.block.AAttn

AAttn(dim, num_heads, area=1)

Bases: Module

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

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

Attributes:

Name Type Description
area int

Number of areas the feature map is divided.

num_heads int

Number of heads into which the attention mechanism is divided.

head_dim int

Dimension of each attention head.

qkv Conv

Convolution layer for computing query, key and value tensors.

proj Conv

Projection convolution layer.

pe Conv

Position encoding convolution layer.

Methods:

Name Description
forward

Applies area-attention to input tensor.

Examples:

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

Parameters:

Name Type Description Default
dim int

Number of hidden channels.

required
num_heads int

Number of heads into which the attention mechanism is divided.

required
area int

Number of areas the feature map is divided, default is 1.

1
Source code in ultralytics/nn/modules/block.py
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
def __init__(self, dim, num_heads, area=1):
    """
    Initialize an Area-attention module for YOLO models.

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

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

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

forward

forward(x)

Process the input tensor through the area-attention.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

required

Returns:

Type Description
Tensor

Output tensor after area-attention.

Source code in ultralytics/nn/modules/block.py
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
def forward(self, x):
    """
    Process the input tensor through the area-attention.

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

    Returns:
        (torch.Tensor): Output tensor after area-attention.
    """
    B, C, H, W = x.shape
    N = H * W

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

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

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

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





ultralytics.nn.modules.block.ABlock

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

Bases: Module

Area-attention block module for efficient feature extraction in YOLO models.

This module implements an area-attention mechanism combined with a feed-forward network for processing feature maps. It uses a novel area-based attention approach that is more efficient than traditional self-attention while maintaining effectiveness.

Attributes:

Name Type Description
attn AAttn

Area-attention module for processing spatial features.

mlp Sequential

Multi-layer perceptron for feature transformation.

Methods:

Name Description
_init_weights

Initializes module weights using truncated normal distribution.

forward

Applies area-attention and feed-forward processing to input tensor.

Examples:

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

Parameters:

Name Type Description Default
dim int

Number of input channels.

required
num_heads int

Number of heads into which the attention mechanism is divided.

required
mlp_ratio float

Expansion ratio for MLP hidden dimension.

1.2
area int

Number of areas the feature map is divided.

1
Source code in ultralytics/nn/modules/block.py
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
def __init__(self, dim, num_heads, mlp_ratio=1.2, area=1):
    """
    Initialize an Area-attention block module.

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

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

    self.apply(self._init_weights)

forward

forward(x)

Forward pass through ABlock.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

required

Returns:

Type Description
Tensor

Output tensor after area-attention and feed-forward processing.

Source code in ultralytics/nn/modules/block.py
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
def forward(self, x):
    """
    Forward pass through ABlock.

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

    Returns:
        (torch.Tensor): Output tensor after area-attention and feed-forward processing.
    """
    x = x + self.attn(x)
    return x + self.mlp(x)





ultralytics.nn.modules.block.A2C2f

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

Bases: Module

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

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

Attributes:

Name Type Description
cv1 Conv

Initial 1x1 convolution layer that reduces input channels to hidden channels.

cv2 Conv

Final 1x1 convolution layer that processes concatenated features.

gamma Parameter | None

Learnable parameter for residual scaling when using area attention.

m ModuleList

List of either ABlock or C3k modules for feature processing.

Methods:

Name Description
forward

Processes input through area-attention or standard convolution pathway.

Examples:

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

Parameters:

Name Type Description Default
c1 int

Number of input channels.

required
c2 int

Number of output channels.

required
n int

Number of ABlock or C3k modules to stack.

1
a2 bool

Whether to use area attention blocks. If False, uses C3k blocks instead.

True
area int

Number of areas the feature map is divided.

1
residual bool

Whether to use residual connections with learnable gamma parameter.

False
mlp_ratio float

Expansion ratio for MLP hidden dimension.

2.0
e float

Channel expansion ratio for hidden channels.

0.5
g int

Number of groups for grouped convolutions.

1
shortcut bool

Whether to use shortcut connections in C3k blocks.

True
Source code in ultralytics/nn/modules/block.py
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
def __init__(self, c1, c2, n=1, a2=True, area=1, residual=False, mlp_ratio=2.0, e=0.5, g=1, shortcut=True):
    """
    Initialize Area-Attention C2f module.

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

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

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

forward

forward(x)

Forward pass through A2C2f layer.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

required

Returns:

Type Description
Tensor

Output tensor after processing.

Source code in ultralytics/nn/modules/block.py
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
def forward(self, x):
    """
    Forward pass through A2C2f layer.

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

    Returns:
        (torch.Tensor): Output tensor after processing.
    """
    y = [self.cv1(x)]
    y.extend(m(y[-1]) for m in self.m)
    y = self.cv2(torch.cat(y, 1))
    if self.gamma is not None:
        return x + self.gamma.view(-1, len(self.gamma), 1, 1) * y
    return y





ultralytics.nn.modules.block.SwiGLUFFN

SwiGLUFFN(gc, ec, e=4)

Bases: Module

SwiGLU Feed-Forward Network for transformer-based architectures.

Source code in ultralytics/nn/modules/block.py
1883
1884
1885
1886
1887
def __init__(self, gc, ec, e=4) -> None:
    """Initialize SwiGLU FFN with input dimension, output dimension, and expansion factor."""
    super().__init__()
    self.w12 = nn.Linear(gc, e * ec)
    self.w3 = nn.Linear(e * ec // 2, ec)

forward

forward(x)

Apply SwiGLU transformation to input features.

Source code in ultralytics/nn/modules/block.py
1889
1890
1891
1892
1893
1894
def forward(self, x):
    """Apply SwiGLU transformation to input features."""
    x12 = self.w12(x)
    x1, x2 = x12.chunk(2, dim=-1)
    hidden = F.silu(x1) * x2
    return self.w3(hidden)





ultralytics.nn.modules.block.Residual

Residual(m)

Bases: Module

Residual connection wrapper for neural network modules.

Source code in ultralytics/nn/modules/block.py
1900
1901
1902
1903
1904
1905
1906
1907
def __init__(self, m) -> None:
    """Initialize residual module with the wrapped module."""
    super().__init__()
    self.m = m
    nn.init.zeros_(self.m.w3.bias)
    # For models with l scale, please change the initialization to
    # nn.init.constant_(self.m.w3.weight, 1e-6)
    nn.init.zeros_(self.m.w3.weight)

forward

forward(x)

Apply residual connection to input features.

Source code in ultralytics/nn/modules/block.py
1909
1910
1911
def forward(self, x):
    """Apply residual connection to input features."""
    return x + self.m(x)





ultralytics.nn.modules.block.SAVPE

SAVPE(ch, c3, embed)

Bases: Module

Spatial-Aware Visual Prompt Embedding module for feature enhancement.

Source code in ultralytics/nn/modules/block.py
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
def __init__(self, ch, c3, embed):
    """Initialize SAVPE module with channels, intermediate channels, and embedding dimension."""
    super().__init__()
    self.cv1 = nn.ModuleList(
        nn.Sequential(
            Conv(x, c3, 3), Conv(c3, c3, 3), nn.Upsample(scale_factor=i * 2) if i in {1, 2} else nn.Identity()
        )
        for i, x in enumerate(ch)
    )

    self.cv2 = nn.ModuleList(
        nn.Sequential(Conv(x, c3, 1), nn.Upsample(scale_factor=i * 2) if i in {1, 2} else nn.Identity())
        for i, x in enumerate(ch)
    )

    self.c = 16
    self.cv3 = nn.Conv2d(3 * c3, embed, 1)
    self.cv4 = nn.Conv2d(3 * c3, self.c, 3, padding=1)
    self.cv5 = nn.Conv2d(1, self.c, 3, padding=1)
    self.cv6 = nn.Sequential(Conv(2 * self.c, self.c, 3), nn.Conv2d(self.c, self.c, 3, padding=1))

forward

forward(x, vp)

Process input features and visual prompts to generate enhanced embeddings.

Source code in ultralytics/nn/modules/block.py
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
def forward(self, x, vp):
    """Process input features and visual prompts to generate enhanced embeddings."""
    y = [self.cv2[i](xi) for i, xi in enumerate(x)]
    y = self.cv4(torch.cat(y, dim=1))

    x = [self.cv1[i](xi) for i, xi in enumerate(x)]
    x = self.cv3(torch.cat(x, dim=1))

    B, C, H, W = x.shape

    Q = vp.shape[1]

    x = x.view(B, C, -1)

    y = y.reshape(B, 1, self.c, H, W).expand(-1, Q, -1, -1, -1).reshape(B * Q, self.c, H, W)
    vp = vp.reshape(B, Q, 1, H, W).reshape(B * Q, 1, H, W)

    y = self.cv6(torch.cat((y, self.cv5(vp)), dim=1))

    y = y.reshape(B, Q, self.c, -1)
    vp = vp.reshape(B, Q, 1, -1)

    score = y * vp + torch.logical_not(vp) * torch.finfo(y.dtype).min

    score = F.softmax(score, dim=-1, dtype=torch.float).to(score.dtype)

    aggregated = score.transpose(-2, -3) @ x.reshape(B, self.c, C // self.c, -1).transpose(-1, -2)

    return F.normalize(aggregated.transpose(-2, -3).reshape(B, Q, -1), dim=-1, p=2)





📅 Created 1 year ago ✏️ Updated 19 days ago