Class CABlock

monai/networks/blocks/cablock.py:78–182

Multi-DConv Head Transposed Self-Attention (MDTA): Differs from standard self-attention by operating on feature channels instead of spatial dimensions. Incorporates depth-wise convolutions for local mixing before attention, achieving linear complexity vs quadratic in vanilla attention. B

class CABlock(nn.Module):
    """Multi-DConv Head Transposed Self-Attention (MDTA): Differs from standard self-attention
    by operating on feature channels instead of spatial dimensions. Incorporates depth-wise
    convolutions for local mixing before attention, achieving linear complexity vs quadratic
    in vanilla attention. Based on SW Zamir, et al., 2022 <https://arxiv.org/abs/2111.09881>

    Args:
        spatial_dims: Number of spatial dimensions (2D or 3D)
        dim: Number of input channels
        num_heads: Number of attention heads
        bias: Whether to use bias in convolution layers
        flash_attention: Whether to use flash attention optimization. Defaults to False.

    Raises:
        ValueError: If flash attention is not available in current PyTorch version
        ValueError: If spatial_dims is greater than 3
    """

    def __init__(self, spatial_dims, dim: int, num_heads: int, bias: bool, flash_attention: bool = False):
        super().__init__()
        if flash_attention and not hasattr(F, "scaled_dot_product_attention"):
            raise ValueError("Flash attention not available")
        if spatial_dims > 3:
            raise ValueError(f"Only 2D and 3D inputs are supported. Got spatial_dims={spatial_dims}")
        self.spatial_dims = spatial_dims
        self.num_heads = num_heads
        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
        self.flash_attention = flash_attention

        self.qkv = Convolution(
            spatial_dims=spatial_dims, in_channels=dim, out_channels=dim * 3, kernel_size=1, bias=bias, conv_only=True
        )

        self.qkv_dwconv = Convolution(
            spatial_dims=spatial_dims,
            in_channels=dim * 3,
            out_channels=dim * 3,
            kernel_size=3,
            strides=1,
            padding=1,
            groups=dim * 3,
            bias=bias,
            conv_only=True,
        )

        self.project_out = Convolution(
            spatial_dims=spatial_dims, in_channels=dim, out_channels=dim, kernel_size=1, bias=bias, conv_only=True
        )

        self._attention_fn = self._get_attention_fn()

    def _get_attention_fn(self):
        if self.flash_attention:
            return self._flash_attention
        return self._normal_attention

    def _flash_attention(self, q, k, v):
        """Flash attention implementation using scaled dot-product attention."""