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

def __init__(self, drop_prob=0.0, scale_by_keep=True):
    """Initialize DropPath module for stochastic depth regularization during training."""
    super().__init__()
    self.drop_prob = drop_prob
    self.scale_by_keep = scale_by_keep

def __init__(
    self,
    embed_dim=256,
    kernel_size=4,
    stride=4,
    padding=0,
    total_stride=16,
    activation=nn.GELU,
):
    """Initializes 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))

def __init__(
    self,
    dim,
    kernel_size=7,
    padding=3,
    drop_path=0.0,
    layer_scale_init_value=1e-6,
    use_dwconv=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()

def __init__(self, layer, num_layers, dim=None, input_projection=False):
    """
    Initializes 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)

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:
    """
    Initializes 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)

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:
    """
    Initializes 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),
            )
        )

def __init__(
    self,
    *args,
    rope_theta=10000.0,
    rope_k_repeat=False,
    feat_sizes=(32, 32),  # [w, h] for stride 16 feats at 512 resolution
    **kwargs,
):
    """Initializes 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

def __init__(
    self,
    dim: int,
    dim_out: int,
    num_heads: int,
    q_pool: nn.Module = None,
):
    """Initializes 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)

def __init__(
    self,
    dim: int,
    dim_out: int,
    num_heads: int,
    mlp_ratio: float = 4.0,
    drop_path: float = 0.0,
    norm_layer: Union[nn.Module, str] = "LayerNorm",
    q_stride: Tuple[int, int] = None,
    act_layer: nn.Module = nn.GELU,
    window_size: int = 0,
):
    """Initializes 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)

def __init__(
    self,
    num_pos_feats,
    temperature: int = 10000,
    normalize: bool = True,
    scale: Optional[float] = None,
):
    """Initializes 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 = {}

def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None:
    """Initializes 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

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: Optional[Tuple[int, int]] = None,
) -> None:
    """
    Initializes 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 (Optional[Tuple[int, int]]): 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