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

def __init__(
    self,
    d_model: int = 256,
    dim_feedforward: int = 2048,
    dropout: float = 0.1,
    pos_enc_at_attn: bool = False,
    pos_enc_at_cross_attn_keys: bool = True,
    pos_enc_at_cross_attn_queries: bool = False,
):
    """
    Initialize a memory attention layer with self-attention, cross-attention, and feedforward components.

    Args:
        d_model (int): Dimensionality of the model.
        dim_feedforward (int): Dimensionality of the feedforward network.
        dropout (float): Dropout rate for regularization.
        pos_enc_at_attn (bool): Whether to add positional encoding at attention.
        pos_enc_at_cross_attn_keys (bool): Whether to add positional encoding to cross-attention keys.
        pos_enc_at_cross_attn_queries (bool): Whether to add positional encoding to cross-attention queries.
    """
    super().__init__()
    self.d_model = d_model
    self.dim_feedforward = dim_feedforward
    self.dropout_value = dropout
    self.self_attn = RoPEAttention(embedding_dim=256, num_heads=1, downsample_rate=1)
    self.cross_attn_image = RoPEAttention(
        rope_k_repeat=True,
        embedding_dim=256,
        num_heads=1,
        downsample_rate=1,
        kv_in_dim=64,
    )

    # Implementation of Feedforward model
    self.linear1 = nn.Linear(d_model, dim_feedforward)
    self.dropout = nn.Dropout(dropout)
    self.linear2 = nn.Linear(dim_feedforward, d_model)

    self.norm1 = nn.LayerNorm(d_model)
    self.norm2 = nn.LayerNorm(d_model)
    self.norm3 = nn.LayerNorm(d_model)
    self.dropout1 = nn.Dropout(dropout)
    self.dropout2 = nn.Dropout(dropout)
    self.dropout3 = nn.Dropout(dropout)

    self.activation = nn.ReLU()

    # Where to add pos enc
    self.pos_enc_at_attn = pos_enc_at_attn
    self.pos_enc_at_cross_attn_queries = pos_enc_at_cross_attn_queries
    self.pos_enc_at_cross_attn_keys = pos_enc_at_cross_attn_keys

def __init__(
    self,
    d_model: int,
    pos_enc_at_input: bool,
    layer: nn.Module,
    num_layers: int,
    batch_first: bool = True,  # Do layers expect batch first input?
):
    """
    Initialize MemoryAttention with specified layers and normalization for sequential data processing.

    This class implements a multi-layer attention mechanism that combines self-attention and cross-attention
    for processing sequential data, particularly useful in transformer-like architectures.

    Args:
        d_model (int): The dimension of the model's hidden state.
        pos_enc_at_input (bool): Whether to apply positional encoding at the input.
        layer (nn.Module): The attention layer to be used in the module.
        num_layers (int): The number of attention layers.
        batch_first (bool): Whether the input tensors are in batch-first format.

    Examples:
        >>> d_model = 256
        >>> layer = MemoryAttentionLayer(d_model)
        >>> attention = MemoryAttention(d_model, pos_enc_at_input=True, layer=layer, num_layers=3)
        >>> curr = torch.randn(10, 32, d_model)  # (seq_len, batch_size, d_model)
        >>> memory = torch.randn(20, 32, d_model)  # (mem_len, batch_size, d_model)
        >>> curr_pos = torch.randn(10, 32, d_model)
        >>> memory_pos = torch.randn(20, 32, d_model)
        >>> output = attention(curr, memory, curr_pos, memory_pos)
        >>> print(output.shape)
        torch.Size([10, 32, 256])
    """
    super().__init__()
    self.d_model = d_model
    self.layers = nn.ModuleList([copy.deepcopy(layer) for _ in range(num_layers)])
    self.num_layers = num_layers
    self.norm = nn.LayerNorm(d_model)
    self.pos_enc_at_input = pos_enc_at_input
    self.batch_first = batch_first