Reference for ultralytics/models/sam/modules/transformer.py

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 Two-Way Transformer for simultaneous attention to image and query points.

    Args:
        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. Must divide embedding_dim.
        mlp_dim (int): Internal channel dimension for the MLP block.
        activation (Type[nn.Module]): Activation function to use in the MLP block.
        attention_downsample_rate (int): Downsampling rate for attention mechanism.
    """
    super().__init__()
    self.depth = depth
    self.embedding_dim = embedding_dim
    self.num_heads = num_heads
    self.mlp_dim = mlp_dim
    self.layers = nn.ModuleList()

    for i in range(depth):
        self.layers.append(
            TwoWayAttentionBlock(
                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),
            )
        )

    self.final_attn_token_to_image = Attention(embedding_dim, num_heads, downsample_rate=attention_downsample_rate)
    self.norm_final_attn = nn.LayerNorm(embedding_dim)

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 TwoWayAttentionBlock for simultaneous attention to image and query points.

    This block implements a specialized transformer layer with four main components: self-attention on sparse
    inputs, cross-attention of sparse inputs to dense inputs, MLP block on sparse inputs, and cross-attention
    of dense inputs to sparse inputs.

    Args:
        embedding_dim (int): Channel dimension of the embeddings.
        num_heads (int): Number of attention heads in the attention layers.
        mlp_dim (int): Hidden dimension of the MLP block.
        activation (Type[nn.Module]): Activation function for the MLP block.
        attention_downsample_rate (int): Downsampling rate for the attention mechanism.
        skip_first_layer_pe (bool): Whether to skip positional encoding in the first layer.
    """
    super().__init__()
    self.self_attn = Attention(embedding_dim, num_heads)
    self.norm1 = nn.LayerNorm(embedding_dim)

    self.cross_attn_token_to_image = Attention(embedding_dim, num_heads, downsample_rate=attention_downsample_rate)
    self.norm2 = nn.LayerNorm(embedding_dim)

    self.mlp = MLPBlock(embedding_dim, mlp_dim, activation)
    self.norm3 = nn.LayerNorm(embedding_dim)

    self.norm4 = nn.LayerNorm(embedding_dim)
    self.cross_attn_image_to_token = Attention(embedding_dim, num_heads, downsample_rate=attention_downsample_rate)

    self.skip_first_layer_pe = skip_first_layer_pe

def __init__(
    self,
    embedding_dim: int,
    num_heads: int,
    downsample_rate: int = 1,
    kv_in_dim: int = None,
) -> None:
    """
    Initialize the Attention module with specified dimensions and settings.

    Args:
        embedding_dim (int): Dimensionality of input embeddings.
        num_heads (int): Number of attention heads.
        downsample_rate (int): Factor by which internal dimensions are downsampled.
        kv_in_dim (int | None): Dimensionality of key and value inputs. If None, uses embedding_dim.

    Raises:
        AssertionError: If num_heads does not evenly divide the internal dim (embedding_dim / downsample_rate).
    """
    super().__init__()
    self.embedding_dim = embedding_dim
    self.kv_in_dim = kv_in_dim if kv_in_dim is not None else embedding_dim
    self.internal_dim = embedding_dim // downsample_rate
    self.num_heads = num_heads
    assert self.internal_dim % num_heads == 0, "num_heads must divide embedding_dim."

    self.q_proj = nn.Linear(embedding_dim, self.internal_dim)
    self.k_proj = nn.Linear(self.kv_in_dim, self.internal_dim)
    self.v_proj = nn.Linear(self.kv_in_dim, self.internal_dim)
    self.out_proj = nn.Linear(self.internal_dim, embedding_dim)