Skip to content

Register now for

# Reference for `ultralytics/models/sam/modules/blocks.py`

Note

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

## `` ultralytics.models.sam.modules.blocks.DropPath

    DropPath(drop_prob: float = 0.0, scale_by_keep: bool = True)

Bases: `Module`

Implements stochastic depth regularization for neural networks during training.

Attributes:

Name | Type | Description
--- | --- | ---
`drop_prob` | `float` | Probability of dropping a path during training.
`scale_by_keep` | `bool` | Whether to scale the output by the keep probability.

Methods:

Name | Description
--- | ---
`forward` | Applies stochastic depth to input tensor during training, with optional scaling.

Examples:

    >>> drop_path = DropPath(drop_prob=0.2, scale_by_keep=True)
    >>> x = torch.randn(32, 64, 224, 224)
    >>> output = drop_path(x)

Source code in `ultralytics/models/sam/modules/blocks.py`

    36
    37
    38
    39
    40
    def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
        """Initialize DropPath module for stochastic depth regularization during training."""
        super().__init__()
        self.drop_prob = drop_prob
        self.scale_by_keep = scale_by_keep

### `` forward

    forward(x: Tensor) -> Tensor

Apply stochastic depth to input tensor during training, with optional scaling.

Source code in `ultralytics/models/sam/modules/blocks.py`

    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    def forward(self, x: Tensor) -> Tensor:
        """Apply stochastic depth to input tensor during training, with optional scaling."""
        if self.drop_prob == 0.0 or not self.training:
            return x
        keep_prob = 1 - self.drop_prob
        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
        random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
        if keep_prob > 0.0 and self.scale_by_keep:
            random_tensor.div_(keep_prob)
        return x * random_tensor

## `` ultralytics.models.sam.modules.blocks.MaskDownSampler

    MaskDownSampler(
        embed_dim: int = 256,
        kernel_size: int = 4,
        stride: int = 4,
        padding: int = 0,
        total_stride: int = 16,
        activation: type[Module] = nn.GELU,
    )

Bases: `Module`

A mask downsampling and embedding module for efficient processing of input masks.

This class implements a mask downsampler that progressively reduces the spatial dimensions of input masks while expanding their channel dimensions using convolutional layers, layer normalization, and activation functions.

Attributes:

Name | Type | Description
--- | --- | ---
`encoder` | `Sequential` | A sequential container of convolutional layers, layer normalization, and activation functions for downsampling and embedding masks.

Methods:

Name | Description
--- | ---
`forward` | Downsamples and encodes input mask to embed_dim channels.

Examples:

    >>> mask_downsampler = MaskDownSampler(embed_dim=256, kernel_size=4, stride=4, padding=0, total_stride=16)
    >>> input_mask = torch.randn(1, 1, 256, 256)
    >>> output = mask_downsampler(input_mask)
    >>> print(output.shape)
    torch.Size([1, 256, 16, 16])

Source code in `ultralytics/models/sam/modules/blocks.py`

     77
     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101
    102
    103
    104
    105
    106
    107
    def __init__(
        self,
        embed_dim: int = 256,
        kernel_size: int = 4,
        stride: int = 4,
        padding: int = 0,
        total_stride: int = 16,
        activation: type[nn.Module] = nn.GELU,
    ):
        """Initialize a mask downsampler module for progressive downsampling and channel expansion."""
        super().__init__()
        num_layers = int(math.log2(total_stride) // math.log2(stride))
        assert stride**num_layers == total_stride
        self.encoder = nn.Sequential()
        mask_in_chans, mask_out_chans = 1, 1
        for _ in range(num_layers):
            mask_out_chans = mask_in_chans * (stride**2)
            self.encoder.append(
                nn.Conv2d(
                    mask_in_chans,
                    mask_out_chans,
                    kernel_size=kernel_size,
                    stride=stride,
                    padding=padding,
                )
            )
            self.encoder.append(LayerNorm2d(mask_out_chans))
            self.encoder.append(activation())
            mask_in_chans = mask_out_chans

        self.encoder.append(nn.Conv2d(mask_out_chans, embed_dim, kernel_size=1))

### `` forward

    forward(x: Tensor) -> Tensor

Downsample and encode input mask to embed_dim channels using convolutional layers and LayerNorm2d.

Source code in `ultralytics/models/sam/modules/blocks.py`

    109
    110
    111
    def forward(self, x: Tensor) -> Tensor:
        """Downsample and encode input mask to embed_dim channels using convolutional layers and LayerNorm2d."""
        return self.encoder(x)

## `` ultralytics.models.sam.modules.blocks.CXBlock

    CXBlock(
        dim: int,
        kernel_size: int = 7,
        padding: int = 3,
        drop_path: float = 0.0,
        layer_scale_init_value: float = 1e-06,
        use_dwconv: bool = True,
    )

Bases: `Module`

ConvNeXt Block for efficient feature extraction in convolutional neural networks.

This block implements a modified version of the ConvNeXt architecture, offering improved performance and flexibility in feature extraction.

Attributes:

Name | Type | Description
--- | --- | ---
`dwconv` | `Conv2d` | Depthwise or standard 2D convolution layer.
`norm` | `LayerNorm2d` | Layer normalization applied to channels.
`pwconv1` | `Linear` | First pointwise convolution implemented as a linear layer.
`act` | `GELU` | GELU activation function.
`pwconv2` | `Linear` | Second pointwise convolution implemented as a linear layer.
`gamma` | `Parameter | None` | Learnable scale parameter for layer scaling.
`drop_path` | `Module` | DropPath layer for stochastic depth regularization.

Methods:

Name | Description
--- | ---
`forward` | Processes the input tensor through the ConvNeXt block.

Examples:

    >>> import torch
    >>> x = torch.randn(1, 64, 56, 56)
    >>> block = CXBlock(dim=64, kernel_size=7, padding=3)
    >>> output = block(x)
    >>> print(output.shape)
    torch.Size([1, 64, 56, 56])

This block implements a modified version of the ConvNeXt architecture, offering improved performance and flexibility in feature extraction.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`dim` | `int` | Number of input channels. | required
`kernel_size` | `int` | Size of the convolutional kernel. | `7`
`padding` | `int` | Padding size for the convolution. | `3`
`drop_path` | `float` | Stochastic depth rate. | `0.0`
`layer_scale_init_value` | `float` | Initial value for Layer Scale. | `1e-06`
`use_dwconv` | `bool` | Whether to use depthwise convolution. | `True`

Examples:

    >>> block = CXBlock(dim=64, kernel_size=7, padding=3)
    >>> x = torch.randn(1, 64, 32, 32)
    >>> output = block(x)
    >>> print(output.shape)
    torch.Size([1, 64, 32, 32])

Source code in `ultralytics/models/sam/modules/blocks.py`

    142
    143
    144
    145
    146
    147
    148
    149
    150
    151
    152
    153
    154
    155
    156
    157
    158
    159
    160
    161
    162
    163
    164
    165
    166
    167
    168
    169
    170
    171
    172
    173
    174
    175
    176
    177
    178
    179
    180
    181
    182
    183
    184
    185
    186
    187
    188
    189
    def __init__(
        self,
        dim: int,
        kernel_size: int = 7,
        padding: int = 3,
        drop_path: float = 0.0,
        layer_scale_init_value: float = 1e-6,
        use_dwconv: bool = True,
    ):
        """
        Initialize a ConvNeXt Block for efficient feature extraction in convolutional neural networks.

        This block implements a modified version of the ConvNeXt architecture, offering improved performance and
        flexibility in feature extraction.

        Args:
            dim (int): Number of input channels.
            kernel_size (int): Size of the convolutional kernel.
            padding (int): Padding size for the convolution.
            drop_path (float): Stochastic depth rate.
            layer_scale_init_value (float): Initial value for Layer Scale.
            use_dwconv (bool): Whether to use depthwise convolution.

        Examples:
            >>> block = CXBlock(dim=64, kernel_size=7, padding=3)
            >>> x = torch.randn(1, 64, 32, 32)
            >>> output = block(x)
            >>> print(output.shape)
            torch.Size([1, 64, 32, 32])
        """
        super().__init__()
        self.dwconv = nn.Conv2d(
            dim,
            dim,
            kernel_size=kernel_size,
            padding=padding,
            groups=dim if use_dwconv else 1,
        )  # depthwise conv
        self.norm = LayerNorm2d(dim, eps=1e-6)
        self.pwconv1 = nn.Linear(dim, 4 * dim)  # pointwise/1x1 convs, implemented with linear layers
        self.act = nn.GELU()
        self.pwconv2 = nn.Linear(4 * dim, dim)
        self.gamma = (
            nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
            if layer_scale_init_value > 0
            else None
        )
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()

### `` forward

    forward(x: Tensor) -> Tensor

Apply ConvNeXt block operations to input tensor, including convolutions and residual connection.

Source code in `ultralytics/models/sam/modules/blocks.py`

    191
    192
    193
    194
    195
    196
    197
    198
    199
    200
    201
    202
    203
    204
    205
    def forward(self, x: Tensor) -> Tensor:
        """Apply ConvNeXt block operations to input tensor, including convolutions and residual connection."""
        input = x
        x = self.dwconv(x)
        x = self.norm(x)
        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        if self.gamma is not None:
            x = self.gamma * x
        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)

        x = input + self.drop_path(x)
        return x

## `` ultralytics.models.sam.modules.blocks.Fuser

    Fuser(
        layer: Module,
        num_layers: int,
        dim: int | None = None,
        input_projection: bool = False,
    )

Bases: `Module`

A module for fusing features through multiple layers of a neural network.

This class applies a series of identical layers to an input tensor, optionally projecting the input first.

Attributes:

Name | Type | Description
--- | --- | ---
`proj` | `Module` | An optional input projection layer. Identity if no projection is needed.
`layers` | `ModuleList` | A list of identical layers to be applied sequentially.

Methods:

Name | Description
--- | ---
`forward` | Applies the fuser to an input tensor.

Examples:

    >>> layer = CXBlock(dim=256)
    >>> fuser = Fuser(layer, num_layers=3, dim=256, input_projection=True)
    >>> x = torch.randn(1, 256, 32, 32)
    >>> output = fuser(x)
    >>> print(output.shape)
    torch.Size([1, 256, 32, 32])

This module creates a sequence of identical layers and optionally applies an input projection.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`layer` | `Module` | The layer to be replicated in the fuser. | required
`num_layers` | `int` | The number of times to replicate the layer. | required
`dim` | `int | None` | The dimension for input projection, if used. | `None`
`input_projection` | `bool` | Whether to use input projection. | `False`

Examples:

    >>> layer = nn.Linear(64, 64)
    >>> fuser = Fuser(layer, num_layers=3, dim=64, input_projection=True)
    >>> input_tensor = torch.randn(1, 64)
    >>> output = fuser(input_tensor)

Source code in `ultralytics/models/sam/modules/blocks.py`

    230
    231
    232
    233
    234
    235
    236
    237
    238
    239
    240
    241
    242
    243
    244
    245
    246
    247
    248
    249
    250
    251
    252
    253
    254
    def __init__(self, layer: nn.Module, num_layers: int, dim: int | None = None, input_projection: bool = False):
        """
        Initialize the Fuser module for feature fusion through multiple layers.

        This module creates a sequence of identical layers and optionally applies an input projection.

        Args:
            layer (nn.Module): The layer to be replicated in the fuser.
            num_layers (int): The number of times to replicate the layer.
            dim (int | None): The dimension for input projection, if used.
            input_projection (bool): Whether to use input projection.

        Examples:
            >>> layer = nn.Linear(64, 64)
            >>> fuser = Fuser(layer, num_layers=3, dim=64, input_projection=True)
            >>> input_tensor = torch.randn(1, 64)
            >>> output = fuser(input_tensor)
        """
        super().__init__()
        self.proj = nn.Identity()
        self.layers = nn.ModuleList([copy.deepcopy(layer) for _ in range(num_layers)])

        if input_projection:
            assert dim is not None
            self.proj = nn.Conv2d(dim, dim, kernel_size=1)

### `` forward

    forward(x: Tensor) -> Tensor

Apply a series of layers to the input tensor, optionally projecting it first.

Source code in `ultralytics/models/sam/modules/blocks.py`

    256
    257
    258
    259
    260
    261
    def forward(self, x: Tensor) -> Tensor:
        """Apply a series of layers to the input tensor, optionally projecting it first."""
        x = self.proj(x)
        for layer in self.layers:
            x = layer(x)
        return x

## `` ultralytics.models.sam.modules.blocks.SAM2TwoWayAttentionBlock

    SAM2TwoWayAttentionBlock(
        embedding_dim: int,
        num_heads: int,
        mlp_dim: int = 2048,
        activation: type[Module] = nn.ReLU,
        attention_downsample_rate: int = 2,
        skip_first_layer_pe: bool = False,
    )

Bases: `TwoWayAttentionBlock`

A two-way attention block for performing self-attention and cross-attention in both directions.

This block extends the TwoWayAttentionBlock and consists of four main components: self-attention on sparse inputs, cross-attention from sparse to dense inputs, an MLP block on sparse inputs, and cross-attention from dense to sparse inputs.

Attributes:

Name | Type | Description
--- | --- | ---
`self_attn` | `Attention` | Self-attention layer for queries.
`norm1` | `LayerNorm` | Layer normalization after the first attention block.
`cross_attn_token_to_image` | `Attention` | Cross-attention layer from queries to keys.
`norm2` | `LayerNorm` | Layer normalization after the second attention block.
`mlp` | `MLP` | MLP block for transforming query embeddings.
`norm3` | `LayerNorm` | Layer normalization after the MLP block.
`norm4` | `LayerNorm` | Layer normalization after the third attention block.
`cross_attn_image_to_token` | `Attention` | Cross-attention layer from keys to queries.
`skip_first_layer_pe` | `bool` | Flag to skip positional encoding in the first layer.

Methods:

Name | Description
--- | ---
`forward` | Processes input through the attention blocks and MLP.

Examples:

    >>> block = SAM2TwoWayAttentionBlock(embedding_dim=256, num_heads=8)
    >>> sparse_input = torch.randn(1, 100, 256)
    >>> dense_input = torch.randn(1, 256, 16, 16)
    >>> sparse_output, dense_output = block(sparse_input, dense_input)

This block extends the TwoWayAttentionBlock and consists of four main components: self-attention on sparse inputs, cross-attention from sparse to dense inputs, an MLP block on sparse inputs, and cross-attention from dense to sparse inputs.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`embedding_dim` | `int` | The channel dimension of the embeddings. | required
`num_heads` | `int` | The number of heads in the attention layers. | required
`mlp_dim` | `int` | The hidden dimension of the MLP block. | `2048`
`activation` | `Type[Module]` | The activation function of the MLP block. | `ReLU`
`attention_downsample_rate` | `int` | The downsample rate for attention computations. | `2`
`skip_first_layer_pe` | `bool` | Whether to skip the positional encoding in the first layer. | `False`

Examples:

    >>> block = SAM2TwoWayAttentionBlock(embedding_dim=256, num_heads=8, mlp_dim=2048)
    >>> sparse_inputs = torch.randn(1, 100, 256)
    >>> dense_inputs = torch.randn(1, 256, 32, 32)
    >>> sparse_outputs, dense_outputs = block(sparse_inputs, dense_inputs)

Source code in `ultralytics/models/sam/modules/blocks.py`

    293
    294
    295
    296
    297
    298
    299
    300
    301
    302
    303
    304
    305
    306
    307
    308
    309
    310
    311
    312
    313
    314
    315
    316
    317
    318
    319
    320
    321
    322
    323
    324
    def __init__(
        self,
        embedding_dim: int,
        num_heads: int,
        mlp_dim: int = 2048,
        activation: type[nn.Module] = nn.ReLU,
        attention_downsample_rate: int = 2,
        skip_first_layer_pe: bool = False,
    ) -> None:
        """
        Initialize a SAM2TwoWayAttentionBlock for performing self-attention and cross-attention in two directions.

        This block extends the TwoWayAttentionBlock and consists of four main components: self-attention on sparse
        inputs, cross-attention from sparse to dense inputs, an MLP block on sparse inputs, and cross-attention
        from dense to sparse inputs.

        Args:
            embedding_dim (int): The channel dimension of the embeddings.
            num_heads (int): The number of heads in the attention layers.
            mlp_dim (int): The hidden dimension of the MLP block.
            activation (Type[nn.Module]): The activation function of the MLP block.
            attention_downsample_rate (int): The downsample rate for attention computations.
            skip_first_layer_pe (bool): Whether to skip the positional encoding in the first layer.

        Examples:
            >>> block = SAM2TwoWayAttentionBlock(embedding_dim=256, num_heads=8, mlp_dim=2048)
            >>> sparse_inputs = torch.randn(1, 100, 256)
            >>> dense_inputs = torch.randn(1, 256, 32, 32)
            >>> sparse_outputs, dense_outputs = block(sparse_inputs, dense_inputs)
        """
        super().__init__(embedding_dim, num_heads, mlp_dim, activation, attention_downsample_rate, skip_first_layer_pe)
        self.mlp = MLP(embedding_dim, mlp_dim, embedding_dim, num_layers=2, act=activation)

## `` ultralytics.models.sam.modules.blocks.SAM2TwoWayTransformer

    SAM2TwoWayTransformer(
        depth: int,
        embedding_dim: int,
        num_heads: int,
        mlp_dim: int,
        activation: type[Module] = nn.ReLU,
        attention_downsample_rate: int = 2,
    )

Bases: `TwoWayTransformer`

A Two-Way Transformer module for simultaneous attention to image and query points.

This class extends the TwoWayTransformer, implementing a specialized transformer decoder that attends to an input image using queries with supplied positional embeddings. It is particularly useful for tasks like object detection, image segmentation, and point cloud processing.

Attributes:

Name | Type | Description
--- | --- | ---
`depth` | `int` | Number of layers in the transformer.
`embedding_dim` | `int` | Channel dimension for input embeddings.
`num_heads` | `int` | Number of heads for multihead attention.
`mlp_dim` | `int` | Internal channel dimension for the MLP block.
`layers` | `ModuleList` | List of SAM2TwoWayAttentionBlock layers comprising the transformer.
`final_attn_token_to_image` | `Attention` | Final attention layer from queries to image.
`norm_final_attn` | `LayerNorm` | Layer normalization applied to final queries.

Methods:

Name | Description
--- | ---
`forward` | Processes input image embeddings and query embeddings through the transformer.

Examples:

    >>> transformer = SAM2TwoWayTransformer(depth=5, embedding_dim=256, num_heads=8, mlp_dim=2048)
    >>> image_embedding = torch.randn(1, 256, 64, 64)
    >>> query_embedding = torch.randn(1, 100, 256)
    >>> output = transformer(image_embedding, query_embedding)
    >>> print(output[0].shape, output[1].shape)
    torch.Size([1, 100, 256]) torch.Size([1, 256, 64, 64])

This transformer decoder attends to an input image using queries with supplied positional embeddings. It is designed for tasks like object detection, image segmentation, and point cloud processing.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`depth` | `int` | Number of layers in the transformer. | required
`embedding_dim` | `int` | Channel dimension for the input embeddings. | required
`num_heads` | `int` | Number of heads for multihead attention. Must divide embedding_dim. | required
`mlp_dim` | `int` | Channel dimension internal to the MLP block. | required
`activation` | `Type[Module]` | Activation function to use in the MLP block. | `ReLU`
`attention_downsample_rate` | `int` | Downsampling rate for attention computations. | `2`

Examples:

    >>> transformer = SAM2TwoWayTransformer(depth=5, embedding_dim=256, num_heads=8, mlp_dim=2048)
    >>> transformer
    SAM2TwoWayTransformer(
      (layers): ModuleList(
        (0-4): 5 x SAM2TwoWayAttentionBlock(...)
      )
      (final_attn_token_to_image): Attention(...)
      (norm_final_attn): LayerNorm(...)
    )

Source code in `ultralytics/models/sam/modules/blocks.py`

    356
    357
    358
    359
    360
    361
    362
    363
    364
    365
    366
    367
    368
    369
    370
    371
    372
    373
    374
    375
    376
    377
    378
    379
    380
    381
    382
    383
    384
    385
    386
    387
    388
    389
    390
    391
    392
    393
    394
    395
    396
    397
    398
    399
    400
    401
    402
    def __init__(
        self,
        depth: int,
        embedding_dim: int,
        num_heads: int,
        mlp_dim: int,
        activation: type[nn.Module] = nn.ReLU,
        attention_downsample_rate: int = 2,
    ) -> None:
        """
        Initialize a SAM2TwoWayTransformer instance.

        This transformer decoder attends to an input image using queries with supplied positional embeddings.
        It is designed for tasks like object detection, image segmentation, and point cloud processing.

        Args:
            depth (int): Number of layers in the transformer.
            embedding_dim (int): Channel dimension for the input embeddings.
            num_heads (int): Number of heads for multihead attention. Must divide embedding_dim.
            mlp_dim (int): Channel dimension internal to the MLP block.
            activation (Type[nn.Module]): Activation function to use in the MLP block.
            attention_downsample_rate (int): Downsampling rate for attention computations.

        Examples:
            >>> transformer = SAM2TwoWayTransformer(depth=5, embedding_dim=256, num_heads=8, mlp_dim=2048)
            >>> transformer
            SAM2TwoWayTransformer(
              (layers): ModuleList(
                (0-4): 5 x SAM2TwoWayAttentionBlock(...)
              )
              (final_attn_token_to_image): Attention(...)
              (norm_final_attn): LayerNorm(...)
            )
        """
        super().__init__(depth, embedding_dim, num_heads, mlp_dim, activation, attention_downsample_rate)
        self.layers = nn.ModuleList()
        for i in range(depth):
            self.layers.append(
                SAM2TwoWayAttentionBlock(
                    embedding_dim=embedding_dim,
                    num_heads=num_heads,
                    mlp_dim=mlp_dim,
                    activation=activation,
                    attention_downsample_rate=attention_downsample_rate,
                    skip_first_layer_pe=(i == 0),
                )
            )

## `` ultralytics.models.sam.modules.blocks.RoPEAttention

    RoPEAttention(
        *args,
        rope_theta: float = 10000.0,
        rope_k_repeat: bool = False,
        feat_sizes: tuple[int, int] = (32, 32),
        **kwargs
    )

Bases: `Attention`

Implements rotary position encoding for attention mechanisms in transformer architectures.

This class extends the base Attention class by incorporating Rotary Position Encoding (RoPE) to enhance the positional awareness of the attention mechanism.

Attributes:

Name | Type | Description
--- | --- | ---
`compute_cis` | `Callable` | Function to compute axial complex numbers for rotary encoding.
`freqs_cis` | `Tensor` | Precomputed frequency tensor for rotary encoding.
`rope_k_repeat` | `bool` | Flag to repeat query RoPE to match key length for cross-attention to memories.

Methods:

Name | Description
--- | ---
`forward` | Applies rotary position encoding and computes attention between query, key, and value tensors.

Examples:

    >>> rope_attn = RoPEAttention(embedding_dim=256, num_heads=8, rope_theta=10000.0, feat_sizes=(32, 32))
    >>> q = torch.randn(1, 1024, 256)
    >>> k = torch.randn(1, 1024, 256)
    >>> v = torch.randn(1, 1024, 256)
    >>> output = rope_attn(q, k, v)
    >>> print(output.shape)
    torch.Size([1, 1024, 256])

Source code in `ultralytics/models/sam/modules/blocks.py`

    430
    431
    432
    433
    434
    435
    436
    437
    438
    439
    440
    441
    442
    443
    444
    def __init__(
        self,
        *args,
        rope_theta: float = 10000.0,
        rope_k_repeat: bool = False,
        feat_sizes: tuple[int, int] = (32, 32),  # [w, h] for stride 16 feats at 512 resolution
        **kwargs,
    ):
        """Initialize RoPEAttention with rotary position encoding for enhanced positional awareness."""
        super().__init__(*args, **kwargs)

        self.compute_cis = partial(compute_axial_cis, dim=self.internal_dim // self.num_heads, theta=rope_theta)
        freqs_cis = self.compute_cis(end_x=feat_sizes[0], end_y=feat_sizes[1])
        self.freqs_cis = freqs_cis
        self.rope_k_repeat = rope_k_repeat  # repeat q rope to match k length, needed for cross-attention to memories

### `` forward

    forward(
        q: Tensor, k: Tensor, v: Tensor, num_k_exclude_rope: int = 0
    ) -> torch.Tensor

Apply rotary position encoding and compute attention between query, key, and value tensors.

Source code in `ultralytics/models/sam/modules/blocks.py`

    446
    447
    448
    449
    450
    451
    452
    453
    454
    455
    456
    457
    458
    459
    460
    461
    462
    463
    464
    465
    466
    467
    468
    469
    470
    471
    472
    473
    474
    475
    476
    477
    478
    479
    480
    481
    482
    483
    484
    485
    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_k_exclude_rope: int = 0) -> torch.Tensor:
        """Apply rotary position encoding and compute attention between query, key, and value tensors."""
        q = self.q_proj(q)
        k = self.k_proj(k)
        v = self.v_proj(v)

        # Separate into heads
        q = self._separate_heads(q, self.num_heads)
        k = self._separate_heads(k, self.num_heads)
        v = self._separate_heads(v, self.num_heads)

        # Apply rotary position encoding
        w = h = math.sqrt(q.shape[-2])
        self.freqs_cis = self.freqs_cis.to(q.device)
        if self.freqs_cis.shape[0] != q.shape[-2]:
            self.freqs_cis = self.compute_cis(end_x=w, end_y=h).to(q.device)
        if q.shape[-2] != k.shape[-2]:
            assert self.rope_k_repeat

        num_k_rope = k.size(-2) - num_k_exclude_rope
        q, k[:, :, :num_k_rope] = apply_rotary_enc(
            q,
            k[:, :, :num_k_rope],
            freqs_cis=self.freqs_cis,
            repeat_freqs_k=self.rope_k_repeat,
        )

        # Attention
        _, _, _, c_per_head = q.shape
        attn = q @ k.permute(0, 1, 3, 2)  # B x N_heads x N_tokens x N_tokens
        attn = attn / math.sqrt(c_per_head)
        attn = torch.softmax(attn, dim=-1)

        # Get output
        out = attn @ v

        out = self._recombine_heads(out)
        out = self.out_proj(out)

        return out

## `` ultralytics.models.sam.modules.blocks.MultiScaleAttention

    MultiScaleAttention(
        dim: int, dim_out: int, num_heads: int, q_pool: Module = None
    )

Bases: `Module`

Implements multiscale self-attention with optional query pooling for efficient feature extraction.

This class provides a flexible implementation of multiscale attention, allowing for optional downsampling of query features through pooling. It's designed to enhance the model's ability to capture multiscale information in visual tasks.

Attributes:

Name | Type | Description
--- | --- | ---
`dim` | `int` | Input dimension of the feature map.
`dim_out` | `int` | Output dimension of the attention module.
`num_heads` | `int` | Number of attention heads.
`scale` | `float` | Scaling factor for dot-product attention.
`q_pool` | `Module | None` | Optional pooling module for query features.
`qkv` | `Linear` | Linear projection for query, key, and value.
`proj` | `Linear` | Output projection.

Methods:

Name | Description
--- | ---
`forward` | Applies multiscale attention to the input tensor.

Examples:

    >>> import torch
    >>> from torch import nn
    >>> x = torch.randn(1, 64, 64, 256)
    >>> msa = MultiScaleAttention(dim=256, dim_out=256, num_heads=8)
    >>> output = msa(x)
    >>> print(output.shape)
    torch.Size([1, 64, 64, 256])

Source code in `ultralytics/models/sam/modules/blocks.py`

    533
    534
    535
    536
    537
    538
    539
    540
    541
    542
    543
    544
    545
    546
    547
    548
    549
    550
    551
    552
    def __init__(
        self,
        dim: int,
        dim_out: int,
        num_heads: int,
        q_pool: nn.Module = None,
    ):
        """Initialize multiscale attention with optional query pooling for efficient feature extraction."""
        super().__init__()

        self.dim = dim
        self.dim_out = dim_out

        self.num_heads = num_heads
        head_dim = dim_out // num_heads
        self.scale = head_dim**-0.5

        self.q_pool = q_pool
        self.qkv = nn.Linear(dim, dim_out * 3)
        self.proj = nn.Linear(dim_out, dim_out)

### `` forward

    forward(x: Tensor) -> torch.Tensor

Apply multiscale attention with optional query pooling to extract multiscale features.

Source code in `ultralytics/models/sam/modules/blocks.py`

    554
    555
    556
    557
    558
    559
    560
    561
    562
    563
    564
    565
    566
    567
    568
    569
    570
    571
    572
    573
    574
    575
    576
    577
    578
    579
    580
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply multiscale attention with optional query pooling to extract multiscale features."""
        B, H, W, _ = x.shape
        # qkv with shape (B, H * W, 3, nHead, C)
        qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1)
        # q, k, v with shape (B, H * W, nheads, C)
        q, k, v = torch.unbind(qkv, 2)

        # Q pooling (for downsample at stage changes)
        if self.q_pool:
            q = do_pool(q.reshape(B, H, W, -1), self.q_pool)
            H, W = q.shape[1:3]  # downsampled shape
            q = q.reshape(B, H * W, self.num_heads, -1)

        # Torch's SDPA expects [B, nheads, H*W, C] so we transpose
        x = F.scaled_dot_product_attention(
            q.transpose(1, 2),
            k.transpose(1, 2),
            v.transpose(1, 2),
        )
        # Transpose back
        x = x.transpose(1, 2)
        x = x.reshape(B, H, W, -1)

        x = self.proj(x)

        return x

## `` ultralytics.models.sam.modules.blocks.MultiScaleBlock

    MultiScaleBlock(
        dim: int,
        dim_out: int,
        num_heads: int,
        mlp_ratio: float = 4.0,
        drop_path: float = 0.0,
        norm_layer: Module | str = "LayerNorm",
        q_stride: tuple[int, int] = None,
        act_layer: type[Module] = nn.GELU,
        window_size: int = 0,
    )

Bases: `Module`

A multiscale attention block with window partitioning and query pooling for efficient vision transformers.

This class implements a multiscale attention mechanism with optional window partitioning and downsampling, designed for use in vision transformer architectures.

Attributes:

Name | Type | Description
--- | --- | ---
`dim` | `int` | Input dimension of the block.
`dim_out` | `int` | Output dimension of the block.
`norm1` | `Module` | First normalization layer.
`window_size` | `int` | Size of the window for partitioning.
`pool` | `Module | None` | Pooling layer for query downsampling.
`q_stride` | `tuple[int, int] | None` | Stride for query pooling.
`attn` | `MultiScaleAttention` | Multi-scale attention module.
`drop_path` | `Module` | Drop path layer for regularization.
`norm2` | `Module` | Second normalization layer.
`mlp` | `MLP` | Multi-layer perceptron module.
`proj` | `Linear | None` | Projection layer for dimension mismatch.

Methods:

Name | Description
--- | ---
`forward` | Processes input tensor through the multiscale block.

Examples:

    >>> block = MultiScaleBlock(dim=256, dim_out=512, num_heads=8, window_size=7)
    >>> x = torch.randn(1, 56, 56, 256)
    >>> output = block(x)
    >>> print(output.shape)
    torch.Size([1, 28, 28, 512])

Source code in `ultralytics/models/sam/modules/blocks.py`

    614
    615
    616
    617
    618
    619
    620
    621
    622
    623
    624
    625
    626
    627
    628
    629
    630
    631
    632
    633
    634
    635
    636
    637
    638
    639
    640
    641
    642
    643
    644
    645
    646
    647
    648
    649
    650
    651
    652
    653
    654
    655
    656
    657
    658
    659
    660
    def __init__(
        self,
        dim: int,
        dim_out: int,
        num_heads: int,
        mlp_ratio: float = 4.0,
        drop_path: float = 0.0,
        norm_layer: nn.Module | str = "LayerNorm",
        q_stride: tuple[int, int] = None,
        act_layer: type[nn.Module] = nn.GELU,
        window_size: int = 0,
    ):
        """Initialize a multiscale attention block with window partitioning and optional query pooling."""
        super().__init__()

        if isinstance(norm_layer, str):
            norm_layer = partial(getattr(nn, norm_layer), eps=1e-6)

        self.dim = dim
        self.dim_out = dim_out
        self.norm1 = norm_layer(dim)

        self.window_size = window_size

        self.pool, self.q_stride = None, q_stride
        if self.q_stride:
            self.pool = nn.MaxPool2d(kernel_size=q_stride, stride=q_stride, ceil_mode=False)

        self.attn = MultiScaleAttention(
            dim,
            dim_out,
            num_heads=num_heads,
            q_pool=self.pool,
        )
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()

        self.norm2 = norm_layer(dim_out)
        self.mlp = MLP(
            dim_out,
            int(dim_out * mlp_ratio),
            dim_out,
            num_layers=2,
            act=act_layer,
        )

        if dim != dim_out:
            self.proj = nn.Linear(dim, dim_out)

### `` forward

    forward(x: Tensor) -> torch.Tensor

Process input through multiscale attention and MLP, with optional windowing and downsampling.

Source code in `ultralytics/models/sam/modules/blocks.py`

    662
    663
    664
    665
    666
    667
    668
    669
    670
    671
    672
    673
    674
    675
    676
    677
    678
    679
    680
    681
    682
    683
    684
    685
    686
    687
    688
    689
    690
    691
    692
    693
    694
    695
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Process input through multiscale attention and MLP, with optional windowing and downsampling."""
        shortcut = x  # B, H, W, C
        x = self.norm1(x)

        # Skip connection
        if self.dim != self.dim_out:
            shortcut = do_pool(self.proj(x), self.pool)

        # Window partition
        window_size = self.window_size
        if window_size > 0:
            H, W = x.shape[1], x.shape[2]
            x, pad_hw = window_partition(x, window_size)

        # Window Attention + Q Pooling (if stage change)
        x = self.attn(x)
        if self.q_stride:
            # Shapes have changed due to Q pooling
            window_size = self.window_size // self.q_stride[0]
            H, W = shortcut.shape[1:3]

            pad_h = (window_size - H % window_size) % window_size
            pad_w = (window_size - W % window_size) % window_size
            pad_hw = (H + pad_h, W + pad_w)

        # Reverse window partition
        if self.window_size > 0:
            x = window_unpartition(x, window_size, pad_hw, (H, W))

        x = shortcut + self.drop_path(x)
        # MLP
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x

## `` ultralytics.models.sam.modules.blocks.PositionEmbeddingSine

    PositionEmbeddingSine(
        num_pos_feats: int,
        temperature: int = 10000,
        normalize: bool = True,
        scale: float | None = None,
    )

Bases: `Module`

A module for generating sinusoidal positional embeddings for 2D inputs like images.

This class implements sinusoidal position encoding for 2D spatial positions, which can be used in transformer-based models for computer vision tasks.

Attributes:

Name | Type | Description
--- | --- | ---
`num_pos_feats` | `int` | Number of positional features (half of the embedding dimension).
`temperature` | `int` | Temperature parameter for the sinusoidal functions.
`normalize` | `bool` | Whether to normalize the positional embeddings.
`scale` | `float` | Scaling factor for the embeddings when normalize is True.
`cache` | `dict` | Cache for storing precomputed embeddings.

Methods:

Name | Description
--- | ---
`_encode_xy` | Encodes 2D positions using sine and cosine functions.
`encode_boxes` | Encodes box coordinates and dimensions into positional embeddings.
`encode_points` | Encodes 2D point coordinates with sinusoidal positional embeddings.
`forward` | Generates sinusoidal position embeddings for 2D inputs.

Examples:

    >>> pos_emb = PositionEmbeddingSine(num_pos_feats=128)
    >>> x = torch.randn(1, 3, 224, 224)
    >>> embeddings = pos_emb(x)
    >>> print(embeddings.shape)
    torch.Size([1, 256, 224, 224])

Source code in `ultralytics/models/sam/modules/blocks.py`

    726
    727
    728
    729
    730
    731
    732
    733
    734
    735
    736
    737
    738
    739
    740
    741
    742
    743
    744
    745
    def __init__(
        self,
        num_pos_feats: int,
        temperature: int = 10000,
        normalize: bool = True,
        scale: float | None = None,
    ):
        """Initialize sinusoidal position embeddings for 2D image inputs."""
        super().__init__()
        assert num_pos_feats % 2 == 0, "Expecting even model width"
        self.num_pos_feats = num_pos_feats // 2
        self.temperature = temperature
        self.normalize = normalize
        if scale is not None and not normalize:
            raise ValueError("normalize should be True if scale is passed")
        if scale is None:
            scale = 2 * math.pi
        self.scale = scale

        self.cache = {}

### `` encode_boxes

    encode_boxes(x: Tensor, y: Tensor, w: Tensor, h: Tensor) -> torch.Tensor

Encode box coordinates and dimensions into positional embeddings for detection.

Source code in `ultralytics/models/sam/modules/blocks.py`

    762
    763
    764
    765
    766
    @torch.no_grad()
    def encode_boxes(self, x: torch.Tensor, y: torch.Tensor, w: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
        """Encode box coordinates and dimensions into positional embeddings for detection."""
        pos_x, pos_y = self._encode_xy(x, y)
        return torch.cat((pos_y, pos_x, h[:, None], w[:, None]), dim=1)

### `` encode_points

    encode_points(x: Tensor, y: Tensor, labels: Tensor) -> torch.Tensor

Encode 2D points with sinusoidal embeddings and append labels.

Source code in `ultralytics/models/sam/modules/blocks.py`

    770
    771
    772
    773
    774
    775
    776
    777
    @torch.no_grad()
    def encode_points(self, x: torch.Tensor, y: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        """Encode 2D points with sinusoidal embeddings and append labels."""
        (bx, nx), (by, ny), (bl, nl) = x.shape, y.shape, labels.shape
        assert bx == by and nx == ny and bx == bl and nx == nl
        pos_x, pos_y = self._encode_xy(x.flatten(), y.flatten())
        pos_x, pos_y = pos_x.reshape(bx, nx, -1), pos_y.reshape(by, ny, -1)
        return torch.cat((pos_y, pos_x, labels[:, :, None]), dim=2)

### `` forward

    forward(x: Tensor) -> Tensor

Generate sinusoidal position embeddings for 2D inputs like images.

Source code in `ultralytics/models/sam/modules/blocks.py`

    779
    780
    781
    782
    783
    784
    785
    786
    787
    788
    789
    790
    791
    792
    793
    794
    795
    796
    797
    798
    799
    800
    801
    802
    803
    804
    805
    806
    807
    808
    809
    810
    @torch.no_grad()
    def forward(self, x: torch.Tensor) -> Tensor:
        """Generate sinusoidal position embeddings for 2D inputs like images."""
        cache_key = (x.shape[-2], x.shape[-1])
        if cache_key in self.cache:
            return self.cache[cache_key][None].repeat(x.shape[0], 1, 1, 1)
        y_embed = (
            torch.arange(1, x.shape[-2] + 1, dtype=torch.float32, device=x.device)
            .view(1, -1, 1)
            .repeat(x.shape[0], 1, x.shape[-1])
        )
        x_embed = (
            torch.arange(1, x.shape[-1] + 1, dtype=torch.float32, device=x.device)
            .view(1, 1, -1)
            .repeat(x.shape[0], x.shape[-2], 1)
        )

        if self.normalize:
            eps = 1e-6
            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale

        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)

        pos_x = x_embed[:, :, :, None] / dim_t
        pos_y = y_embed[:, :, :, None] / dim_t
        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
        self.cache[cache_key] = pos[0]
        return pos

## `` ultralytics.models.sam.modules.blocks.PositionEmbeddingRandom

    PositionEmbeddingRandom(num_pos_feats: int = 64, scale: float | None = None)

Bases: `Module`

Positional encoding using random spatial frequencies.

This class generates positional embeddings for input coordinates using random spatial frequencies. It is particularly useful for transformer-based models that require position information.

Attributes:

Name | Type | Description
--- | --- | ---
`positional_encoding_gaussian_matrix` | `Tensor` | A buffer containing random values for encoding.

Methods:

Name | Description
--- | ---
`_pe_encoding` | Positionally encodes points that are normalized to [0,1].
`forward` | Generates positional encoding for a grid of the specified size.
`forward_with_coords` | Positionally encodes points that are not normalized to [0,1].

Examples:

    >>> pe = PositionEmbeddingRandom(num_pos_feats=64)
    >>> size = (32, 32)
    >>> encoding = pe(size)
    >>> print(encoding.shape)
    torch.Size([128, 32, 32])

Source code in `ultralytics/models/sam/modules/blocks.py`

    836
    837
    838
    839
    840
    841
    842
    843
    844
    845
    def __init__(self, num_pos_feats: int = 64, scale: float | None = None) -> None:
        """Initialize random spatial frequency position embedding for transformers."""
        super().__init__()
        if scale is None or scale <= 0.0:
            scale = 1.0
        self.register_buffer("positional_encoding_gaussian_matrix", scale * torch.randn((2, num_pos_feats)))

        # Set non-deterministic for forward() error 'cumsum_cuda_kernel does not have a deterministic implementation'
        torch.use_deterministic_algorithms(False)
        torch.backends.cudnn.deterministic = False

### `` forward

    forward(size: tuple[int, int]) -> torch.Tensor

Generate positional encoding for a grid using random spatial frequencies.

Source code in `ultralytics/models/sam/modules/blocks.py`

    856
    857
    858
    859
    860
    861
    862
    863
    864
    865
    866
    867
    868
    869
    870
    def forward(self, size: tuple[int, int]) -> torch.Tensor:
        """Generate positional encoding for a grid using random spatial frequencies."""
        h, w = size
        grid = torch.ones(
            (h, w),
            device=self.positional_encoding_gaussian_matrix.device,
            dtype=self.positional_encoding_gaussian_matrix.dtype,
        )
        y_embed = grid.cumsum(dim=0) - 0.5
        x_embed = grid.cumsum(dim=1) - 0.5
        y_embed = y_embed / h
        x_embed = x_embed / w

        pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))
        return pe.permute(2, 0, 1)  # C x H x W

### `` forward_with_coords

    forward_with_coords(
        coords_input: Tensor, image_size: tuple[int, int]
    ) -> torch.Tensor

Positionally encode input coordinates, normalizing them to [0,1] based on the given image size.

Source code in `ultralytics/models/sam/modules/blocks.py`

    872
    873
    874
    875
    876
    877
    def forward_with_coords(self, coords_input: torch.Tensor, image_size: tuple[int, int]) -> torch.Tensor:
        """Positionally encode input coordinates, normalizing them to [0,1] based on the given image size."""
        coords = coords_input.clone()
        coords[:, :, 0] = coords[:, :, 0] / image_size[1]
        coords[:, :, 1] = coords[:, :, 1] / image_size[0]
        return self._pe_encoding(coords)  # B x N x C

## `` ultralytics.models.sam.modules.blocks.Block

    Block(
        dim: int,
        num_heads: int,
        mlp_ratio: float = 4.0,
        qkv_bias: bool = True,
        norm_layer: type[Module] = nn.LayerNorm,
        act_layer: type[Module] = nn.GELU,
        use_rel_pos: bool = False,
        rel_pos_zero_init: bool = True,
        window_size: int = 0,
        input_size: tuple[int, int] | None = None,
    )

Bases: `Module`

Transformer block with support for window attention and residual propagation.

This class implements a transformer block that can use either global or windowed self-attention, followed by a feed-forward network. It supports relative positional embeddings and is designed for use in vision transformer architectures.

Attributes:

Name | Type | Description
--- | --- | ---
`norm1` | `Module` | First normalization layer.
`attn` | `REAttention` | Self-attention layer with optional relative positional encoding.
`norm2` | `Module` | Second normalization layer.
`mlp` | `MLPBlock` | Multi-layer perceptron block.
`window_size` | `int` | Size of attention window. If 0, global attention is used.

Methods:

Name | Description
--- | ---
`forward` | Processes input through the transformer block.

Examples:

    >>> import torch
    >>> block = Block(dim=256, num_heads=8, window_size=7)
    >>> x = torch.randn(1, 56, 56, 256)
    >>> output = block(x)
    >>> print(output.shape)
    torch.Size([1, 56, 56, 256])

This constructor sets up a transformer block that can use either global or windowed self-attention, followed by a feed-forward network. It supports relative positional embeddings and is designed for use in vision transformer architectures.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`dim` | `int` | Number of input channels. | required
`num_heads` | `int` | Number of attention heads in the self-attention layer. | required
`mlp_ratio` | `float` | Ratio of mlp hidden dimension to embedding dimension. | `4.0`
`qkv_bias` | `bool` | If True, adds a learnable bias to query, key, value projections. | `True`
`norm_layer` | `Type[Module]` | Type of normalization layer to use. | `LayerNorm`
`act_layer` | `Type[Module]` | Type of activation function to use in the MLP block. | `GELU`
`use_rel_pos` | `bool` | If True, uses relative positional embeddings in attention. | `False`
`rel_pos_zero_init` | `bool` | If True, initializes relative positional parameters to zero. | `True`
`window_size` | `int` | Size of attention window. If 0, uses global attention. | `0`
`input_size` | `tuple[int, int] | None` | Input resolution for calculating relative positional parameter size. | `None`

Examples:

    >>> block = Block(dim=256, num_heads=8, window_size=7)
    >>> x = torch.randn(1, 56, 56, 256)
    >>> output = block(x)
    >>> print(output.shape)
    torch.Size([1, 56, 56, 256])

Source code in `ultralytics/models/sam/modules/blocks.py`

    907
    908
    909
    910
    911
    912
    913
    914
    915
    916
    917
    918
    919
    920
    921
    922
    923
    924
    925
    926
    927
    928
    929
    930
    931
    932
    933
    934
    935
    936
    937
    938
    939
    940
    941
    942
    943
    944
    945
    946
    947
    948
    949
    950
    951
    952
    953
    954
    955
    956
    957
    958
    959
    960
    def __init__(
        self,
        dim: int,
        num_heads: int,
        mlp_ratio: float = 4.0,
        qkv_bias: bool = True,
        norm_layer: type[nn.Module] = nn.LayerNorm,
        act_layer: type[nn.Module] = nn.GELU,
        use_rel_pos: bool = False,
        rel_pos_zero_init: bool = True,
        window_size: int = 0,
        input_size: tuple[int, int] | None = None,
    ) -> None:
        """
        Initialize a transformer block with optional window attention and relative positional embeddings.

        This constructor sets up a transformer block that can use either global or windowed self-attention,
        followed by a feed-forward network. It supports relative positional embeddings and is designed
        for use in vision transformer architectures.

        Args:
            dim (int): Number of input channels.
            num_heads (int): Number of attention heads in the self-attention layer.
            mlp_ratio (float): Ratio of mlp hidden dimension to embedding dimension.
            qkv_bias (bool): If True, adds a learnable bias to query, key, value projections.
            norm_layer (Type[nn.Module]): Type of normalization layer to use.
            act_layer (Type[nn.Module]): Type of activation function to use in the MLP block.
            use_rel_pos (bool): If True, uses relative positional embeddings in attention.
            rel_pos_zero_init (bool): If True, initializes relative positional parameters to zero.
            window_size (int): Size of attention window. If 0, uses global attention.
            input_size (tuple[int, int] | None): Input resolution for calculating relative positional parameter size.

        Examples:
            >>> block = Block(dim=256, num_heads=8, window_size=7)
            >>> x = torch.randn(1, 56, 56, 256)
            >>> output = block(x)
            >>> print(output.shape)
            torch.Size([1, 56, 56, 256])
        """
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = REAttention(
            dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            use_rel_pos=use_rel_pos,
            rel_pos_zero_init=rel_pos_zero_init,
            input_size=input_size if window_size == 0 else (window_size, window_size),
        )

        self.norm2 = norm_layer(dim)
        self.mlp = MLPBlock(embedding_dim=dim, mlp_dim=int(dim * mlp_ratio), act=act_layer)

        self.window_size = window_size

### `` forward

    forward(x: Tensor) -> torch.Tensor

Process input through transformer block with optional windowed self-attention and residual connection.

Source code in `ultralytics/models/sam/modules/blocks.py`

    962
    963
    964
    965
    966
    967
    968
    969
    970
    971
    972
    973
    974
    975
    976
    977
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Process input through transformer block with optional windowed self-attention and residual connection."""
        shortcut = x
        x = self.norm1(x)
        # Window partition
        if self.window_size > 0:
            H, W = x.shape[1], x.shape[2]
            x, pad_hw = window_partition(x, self.window_size)

        x = self.attn(x)
        # Reverse window partition
        if self.window_size > 0:
            x = window_unpartition(x, self.window_size, pad_hw, (H, W))

        x = shortcut + x
        return x + self.mlp(self.norm2(x))

## `` ultralytics.models.sam.modules.blocks.REAttention

    REAttention(
        dim: int,
        num_heads: int = 8,
        qkv_bias: bool = True,
        use_rel_pos: bool = False,
        rel_pos_zero_init: bool = True,
        input_size: tuple[int, int] | None = None,
    )

Bases: `Module`

Relative Position Attention module for efficient self-attention in transformer architectures.

This class implements a multi-head attention mechanism with relative positional embeddings, designed for use in vision transformer models. It supports optional query pooling and window partitioning for efficient processing of large inputs.

Attributes:

Name | Type | Description
--- | --- | ---
`num_heads` | `int` | Number of attention heads.
`scale` | `float` | Scaling factor for attention computation.
`qkv` | `Linear` | Linear projection for query, key, and value.
`proj` | `Linear` | Output projection layer.
`use_rel_pos` | `bool` | Whether to use relative positional embeddings.
`rel_pos_h` | `Parameter` | Relative positional embeddings for height dimension.
`rel_pos_w` | `Parameter` | Relative positional embeddings for width dimension.

Methods:

Name | Description
--- | ---
`forward` | Applies multi-head attention with optional relative positional encoding to input tensor.

Examples:

    >>> attention = REAttention(dim=256, num_heads=8, input_size=(32, 32))
    >>> x = torch.randn(1, 32, 32, 256)
    >>> output = attention(x)
    >>> print(output.shape)
    torch.Size([1, 32, 32, 256])

This module implements multi-head attention with optional relative positional encodings, designed specifically for vision tasks in transformer models.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`dim` | `int` | Number of input channels. | required
`num_heads` | `int` | Number of attention heads. | `8`
`qkv_bias` | `bool` | If True, adds a learnable bias to query, key, value projections. | `True`
`use_rel_pos` | `bool` | If True, uses relative positional encodings. | `False`
`rel_pos_zero_init` | `bool` | If True, initializes relative positional parameters to zero. | `True`
`input_size` | `tuple[int, int] | None` | Input resolution for calculating relative positional parameter size. Required if use_rel_pos is True. | `None`

Examples:

    >>> attention = REAttention(dim=256, num_heads=8, input_size=(32, 32))
    >>> x = torch.randn(1, 32, 32, 256)
    >>> output = attention(x)
    >>> print(output.shape)
    torch.Size([1, 32, 32, 256])

Source code in `ultralytics/models/sam/modules/blocks.py`

    1008
    1009
    1010
    1011
    1012
    1013
    1014
    1015
    1016
    1017
    1018
    1019
    1020
    1021
    1022
    1023
    1024
    1025
    1026
    1027
    1028
    1029
    1030
    1031
    1032
    1033
    1034
    1035
    1036
    1037
    1038
    1039
    1040
    1041
    1042
    1043
    1044
    1045
    1046
    1047
    1048
    1049
    1050
    1051
    1052
    def __init__(
        self,
        dim: int,
        num_heads: int = 8,
        qkv_bias: bool = True,
        use_rel_pos: bool = False,
        rel_pos_zero_init: bool = True,
        input_size: tuple[int, int] | None = None,
    ) -> None:
        """
        Initialize a Relative Position Attention module for transformer-based architectures.

        This module implements multi-head attention with optional relative positional encodings, designed
        specifically for vision tasks in transformer models.

        Args:
            dim (int): Number of input channels.
            num_heads (int): Number of attention heads.
            qkv_bias (bool): If True, adds a learnable bias to query, key, value projections.
            use_rel_pos (bool): If True, uses relative positional encodings.
            rel_pos_zero_init (bool): If True, initializes relative positional parameters to zero.
            input_size (tuple[int, int] | None): Input resolution for calculating relative positional parameter size.
                Required if use_rel_pos is True.

        Examples:
            >>> attention = REAttention(dim=256, num_heads=8, input_size=(32, 32))
            >>> x = torch.randn(1, 32, 32, 256)
            >>> output = attention(x)
            >>> print(output.shape)
            torch.Size([1, 32, 32, 256])
        """
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim**-0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.proj = nn.Linear(dim, dim)

        self.use_rel_pos = use_rel_pos
        if self.use_rel_pos:
            assert input_size is not None, "Input size must be provided if using relative positional encoding."
            # Initialize relative positional embeddings
            self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim))
            self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim))

### `` forward

    forward(x: Tensor) -> torch.Tensor

Apply multi-head attention with optional relative positional encoding to input tensor.

Source code in `ultralytics/models/sam/modules/blocks.py`

    1054
    1055
    1056
    1057
    1058
    1059
    1060
    1061
    1062
    1063
    1064
    1065
    1066
    1067
    1068
    1069
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply multi-head attention with optional relative positional encoding to input tensor."""
        B, H, W, _ = x.shape
        # qkv with shape (3, B, nHead, H * W, C)
        qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
        # q, k, v with shape (B * nHead, H * W, C)
        q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0)

        attn = (q * self.scale) @ k.transpose(-2, -1)

        if self.use_rel_pos:
            attn = add_decomposed_rel_pos(attn, q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W))

        attn = attn.softmax(dim=-1)
        x = (attn @ v).view(B, self.num_heads, H, W, -1).permute(0, 2, 3, 1, 4).reshape(B, H, W, -1)
        return self.proj(x)

## `` ultralytics.models.sam.modules.blocks.PatchEmbed

    PatchEmbed(
        kernel_size: tuple[int, int] = (16, 16),
        stride: tuple[int, int] = (16, 16),
        padding: tuple[int, int] = (0, 0),
        in_chans: int = 3,
        embed_dim: int = 768,
    )

Bases: `Module`

Image to Patch Embedding module for vision transformer architectures.

This module converts an input image into a sequence of patch embeddings using a convolutional layer. It is commonly used as the first layer in vision transformer architectures to transform image data into a suitable format for subsequent transformer blocks.

Attributes:

Name | Type | Description
--- | --- | ---
`proj` | `Conv2d` | Convolutional layer for projecting image patches to embeddings.

Methods:

Name | Description
--- | ---
`forward` | Applies patch embedding to the input tensor.

Examples:

    >>> patch_embed = PatchEmbed(kernel_size=(16, 16), stride=(16, 16), in_chans=3, embed_dim=768)
    >>> x = torch.randn(1, 3, 224, 224)
    >>> output = patch_embed(x)
    >>> print(output.shape)
    torch.Size([1, 768, 14, 14])

This module is typically used as the first layer in vision transformer architectures to transform image data into a suitable format for subsequent transformer blocks.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`kernel_size` | `tuple[int, int]` | Size of the convolutional kernel for patch extraction. | `(16, 16)`
`stride` | `tuple[int, int]` | Stride of the convolutional operation. | `(16, 16)`
`padding` | `tuple[int, int]` | Padding applied to the input before convolution. | `(0, 0)`
`in_chans` | `int` | Number of input image channels. | `3`
`embed_dim` | `int` | Dimensionality of the output patch embeddings. | `768`

Examples:

    >>> patch_embed = PatchEmbed(kernel_size=(16, 16), stride=(16, 16), in_chans=3, embed_dim=768)
    >>> x = torch.randn(1, 3, 224, 224)
    >>> output = patch_embed(x)
    >>> print(output.shape)
    torch.Size([1, 768, 14, 14])

Source code in `ultralytics/models/sam/modules/blocks.py`

    1094
    1095
    1096
    1097
    1098
    1099
    1100
    1101
    1102
    1103
    1104
    1105
    1106
    1107
    1108
    1109
    1110
    1111
    1112
    1113
    1114
    1115
    1116
    1117
    1118
    1119
    1120
    1121
    1122
    1123
    1124
    def __init__(
        self,
        kernel_size: tuple[int, int] = (16, 16),
        stride: tuple[int, int] = (16, 16),
        padding: tuple[int, int] = (0, 0),
        in_chans: int = 3,
        embed_dim: int = 768,
    ) -> None:
        """
        Initialize the PatchEmbed module for converting image patches to embeddings.

        This module is typically used as the first layer in vision transformer architectures to transform
        image data into a suitable format for subsequent transformer blocks.

        Args:
            kernel_size (tuple[int, int]): Size of the convolutional kernel for patch extraction.
            stride (tuple[int, int]): Stride of the convolutional operation.
            padding (tuple[int, int]): Padding applied to the input before convolution.
            in_chans (int): Number of input image channels.
            embed_dim (int): Dimensionality of the output patch embeddings.

        Examples:
            >>> patch_embed = PatchEmbed(kernel_size=(16, 16), stride=(16, 16), in_chans=3, embed_dim=768)
            >>> x = torch.randn(1, 3, 224, 224)
            >>> output = patch_embed(x)
            >>> print(output.shape)
            torch.Size([1, 768, 14, 14])
        """
        super().__init__()

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding)

### `` forward

    forward(x: Tensor) -> torch.Tensor

Compute patch embedding by applying convolution and transposing resulting tensor.

Source code in `ultralytics/models/sam/modules/blocks.py`

    1126
    1127
    1128
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Compute patch embedding by applying convolution and transposing resulting tensor."""
        return self.proj(x).permute(0, 2, 3, 1)  # B C H W -> B H W C

## `` ultralytics.models.sam.modules.blocks.do_pool

    do_pool(x: Tensor, pool: Module, norm: Module = None) -> torch.Tensor

Apply pooling and optional normalization to a tensor, handling spatial dimension permutations.

Source code in `ultralytics/models/sam/modules/blocks.py`

    488
    489
    490
    491
    492
    493
    494
    495
    496
    497
    498
    499
    500
    def do_pool(x: torch.Tensor, pool: nn.Module, norm: nn.Module = None) -> torch.Tensor:
        """Apply pooling and optional normalization to a tensor, handling spatial dimension permutations."""
        if pool is None:
            return x
        # (B, H, W, C) -> (B, C, H, W)
        x = x.permute(0, 3, 1, 2)
        x = pool(x)
        # (B, C, H', W') -> (B, H', W', C)
        x = x.permute(0, 2, 3, 1)
        if norm:
            x = norm(x)

        return x

📅 Created 1 year ago ✏️ Updated 1 year ago