Method __init__

monai/networks/blocks/selfattention.py:32–161

Args: hidden_size (int): dimension of hidden layer. num_heads (int): number of attention heads. dropout_rate (float, optional): fraction of the input units to drop. Defaults to 0.0. qkv_bias (bool, optional): bias term for the qkv linear layer

(
        self,
        hidden_size: int,
        num_heads: int,
        dropout_rate: float = 0.0,
        qkv_bias: bool = False,
        save_attn: bool = False,
        dim_head: int | None = None,
        hidden_input_size: int | None = None,
        causal: bool = False,
        sequence_length: int | None = None,
        rel_pos_embedding: str | None = None,
        input_size: tuple | None = None,
        attention_dtype: torch.dtype | None = None,
        include_fc: bool = True,
        use_combined_linear: bool = True,
        use_flash_attention: bool = False,
    )

    """

    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        dropout_rate: float = 0.0,
        qkv_bias: bool = False,
        save_attn: bool = False,
        dim_head: int | None = None,
        hidden_input_size: int | None = None,
        causal: bool = False,
        sequence_length: int | None = None,
        rel_pos_embedding: str | None = None,
        input_size: tuple | None = None,
        attention_dtype: torch.dtype | None = None,
        include_fc: bool = True,
        use_combined_linear: bool = True,
        use_flash_attention: bool = False,
    ) -> None:
        """
        Args:
            hidden_size (int): dimension of hidden layer.
            num_heads (int): number of attention heads.
            dropout_rate (float, optional): fraction of the input units to drop. Defaults to 0.0.
            qkv_bias (bool, optional): bias term for the qkv linear layer. Defaults to False.
            save_attn (bool, optional): to make accessible the attention matrix. Defaults to False.
            dim_head (int, optional): dimension of each head. Defaults to hidden_size // num_heads.
            hidden_input_size (int, optional): dimension of the input tensor. Defaults to hidden_size.
            causal: whether to use causal attention (see https://arxiv.org/abs/1706.03762).
            sequence_length: if causal is True, it is necessary to specify the sequence length.
            rel_pos_embedding (str, optional): Add relative positional embeddings to the attention map.
                For now only "decomposed" is supported (see https://arxiv.org/abs/2112.01526). 2D and 3D are supported.
            input_size (tuple(spatial_dim), optional): Input resolution for calculating the relative
                positional parameter size.
            attention_dtype: cast attention operations to this dtype.
            include_fc: whether to include the final linear layer. Default to True.
            use_combined_linear: whether to use a single linear layer for qkv projection, default to True.
            use_flash_attention: if True, use Pytorch's inbuilt flash attention for a memory efficient attention mechanism
                (see https://pytorch.org/docs/2.2/generated/torch.nn.functional.scaled_dot_product_attention.html).

        Raises:
            ValueError: if ``dropout_rate`` is not between 0 and 1.
            ValueError: if ``hidden_size`` is not divisible by ``num_heads``.
            ValueError: if ``causal`` is True and ``sequence_length`` is not provided.
            ValueError: if both ``save_attn`` and ``use_flash_attention`` are True.
            ValueError: if ``rel_pos_embedding`` is not None and ``use_flash_attention`` is True.

        """

        super().__init__()

        if not (0 <= dropout_rate <= 1):
            raise ValueError("dropout_rate should be between 0 and 1.")

        if hidden_size % num_heads != 0:
            raise ValueError("hidden size should be divisible by num_heads.")

        if dim_head:
            self.inner_dim = num_heads * dim_head