Class MDTATransformerBlock

monai/networks/nets/restormer.py:24–58

Basic transformer unit combining MDTA and GDFN with skip connections. Unlike standard transformers that use LayerNorm, this block uses Instance Norm for better adaptation to image restoration tasks. Args: spatial_dims: Number of spatial dimensions (2D or 3D) dim: Number

class MDTATransformerBlock(nn.Module):
    """Basic transformer unit combining MDTA and GDFN with skip connections.
    Unlike standard transformers that use LayerNorm, this block uses Instance Norm
    for better adaptation to image restoration tasks.

    Args:
        spatial_dims: Number of spatial dimensions (2D or 3D)
        dim: Number of input channels
        num_heads: Number of attention heads
        ffn_expansion_factor: Expansion factor for feed-forward network
        bias: Whether to use bias in attention layers
        layer_norm_use_bias: Whether to use bias in layer normalization. Defaults to False.
        flash_attention: Whether to use flash attention optimization. Defaults to False.
    """

    def __init__(
        self,
        spatial_dims: int,
        dim: int,
        num_heads: int,
        ffn_expansion_factor: float,
        bias: bool,
        layer_norm_use_bias: bool = False,
        flash_attention: bool = False,
    ):
        super().__init__()
        self.norm1 = Norm[Norm.INSTANCE, spatial_dims](dim, affine=layer_norm_use_bias)
        self.attn = CABlock(spatial_dims, dim, num_heads, bias, flash_attention)
        self.norm2 = Norm[Norm.INSTANCE, spatial_dims](dim, affine=layer_norm_use_bias)
        self.ffn = FeedForward(spatial_dims, dim, ffn_expansion_factor, bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x


class OverlapPatchEmbed(Convolution):