Equimo 0.3.5__tar.gz → 0.4.0a2__tar.gz

{equimo-0.3.5/src/Equimo.egg-info → equimo-0.4.0a2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Equimo
-Version: 0.3.5
+Version: 0.4.0a2
 Summary: Implementation of popular vision models in Jax
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
@@ -76,6 +76,8 @@ Beyond normal ViT (e.g., dinov2 or siglip), equimo proposes other SotA architect
 | PartialFormer | [Efficient Vision Transformers with Partial Attention](https://eccv.ecva.net/virtual/2024/poster/1877)                                                          | 2024 | ✅        |
 | SHViT         | [SHViT: Single-Head Vision Transformer with Memory Efficient Macro Design](https://arxiv.org/abs/2401.16456)                                                    | 2024 | ✅        |
 | VSSD          | [VSSD: Vision Mamba with Non-Causal State Space Duality](https://arxiv.org/abs/2407.18559)                                                                      | 2024 | ✅        |
+| ReduceFormer  | [ReduceFormer: Attention with Tensor Reduction by Summation](https://arxiv.org/abs/2406.07488)                                                                  | 2024 | ✅        |
+| LowFormer     | [LowFormer: Hardware Efficient Design for Convolutional Transformer Backbones](https://arxiv.org/abs/2409.03460)                                                | 2024 | ✅        |
 
 \*: Only contains the [Linear Angular Attention](https://github.com/clementpoiret/Equimo/blob/f8fcc79e45ca65e9deb1d970c4286c0b8562f9c2/equimo/layers/attention.py#L1407) module. It is straight forward to build a ViT around it, but may require an additional `__call__` kwarg to control the `sparse_reg` bool.
 

{equimo-0.3.5 → equimo-0.4.0a2}/README.md

@@ -40,6 +40,8 @@ Beyond normal ViT (e.g., dinov2 or siglip), equimo proposes other SotA architect
 | PartialFormer | [Efficient Vision Transformers with Partial Attention](https://eccv.ecva.net/virtual/2024/poster/1877)                                                          | 2024 | ✅        |
 | SHViT         | [SHViT: Single-Head Vision Transformer with Memory Efficient Macro Design](https://arxiv.org/abs/2401.16456)                                                    | 2024 | ✅        |
 | VSSD          | [VSSD: Vision Mamba with Non-Causal State Space Duality](https://arxiv.org/abs/2407.18559)                                                                      | 2024 | ✅        |
+| ReduceFormer  | [ReduceFormer: Attention with Tensor Reduction by Summation](https://arxiv.org/abs/2406.07488)                                                                  | 2024 | ✅        |
+| LowFormer     | [LowFormer: Hardware Efficient Design for Convolutional Transformer Backbones](https://arxiv.org/abs/2409.03460)                                                | 2024 | ✅        |
 
 \*: Only contains the [Linear Angular Attention](https://github.com/clementpoiret/Equimo/blob/f8fcc79e45ca65e9deb1d970c4286c0b8562f9c2/equimo/layers/attention.py#L1407) module. It is straight forward to build a ViT around it, but may require an additional `__call__` kwarg to control the `sparse_reg` bool.
 

{equimo-0.3.5 → equimo-0.4.0a2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "Equimo"
-version = "0.3.5"
+version = "0.4.0-alpha.2"
 description = "Implementation of popular vision models in Jax"
 readme = "README.md"
 requires-python = ">=3.11"

{equimo-0.3.5 → equimo-0.4.0a2/src/Equimo.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Equimo
-Version: 0.3.5
+Version: 0.4.0a2
 Summary: Implementation of popular vision models in Jax
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
@@ -76,6 +76,8 @@ Beyond normal ViT (e.g., dinov2 or siglip), equimo proposes other SotA architect
 | PartialFormer | [Efficient Vision Transformers with Partial Attention](https://eccv.ecva.net/virtual/2024/poster/1877)                                                          | 2024 | ✅        |
 | SHViT         | [SHViT: Single-Head Vision Transformer with Memory Efficient Macro Design](https://arxiv.org/abs/2401.16456)                                                    | 2024 | ✅        |
 | VSSD          | [VSSD: Vision Mamba with Non-Causal State Space Duality](https://arxiv.org/abs/2407.18559)                                                                      | 2024 | ✅        |
+| ReduceFormer  | [ReduceFormer: Attention with Tensor Reduction by Summation](https://arxiv.org/abs/2406.07488)                                                                  | 2024 | ✅        |
+| LowFormer     | [LowFormer: Hardware Efficient Design for Convolutional Transformer Backbones](https://arxiv.org/abs/2409.03460)                                                | 2024 | ✅        |
 
 \*: Only contains the [Linear Angular Attention](https://github.com/clementpoiret/Equimo/blob/f8fcc79e45ca65e9deb1d970c4286c0b8562f9c2/equimo/layers/attention.py#L1407) module. It is straight forward to build a ViT around it, but may require an additional `__call__` kwarg to control the `sparse_reg` bool.
 

{equimo-0.3.5 → equimo-0.4.0a2}/src/Equimo.egg-info/SOURCES.txt

@@ -30,8 +30,10 @@ src/equimo/layers/squeeze_excite.py
 src/equimo/models/__init__.py
 src/equimo/models/emamodel.py
 src/equimo/models/fastervit.py
+src/equimo/models/lowformer.py
 src/equimo/models/mlla.py
 src/equimo/models/partialformer.py
+src/equimo/models/reduceformer.py
 src/equimo/models/shvit.py
 src/equimo/models/vit.py
 src/equimo/models/vssd.py

equimo-0.4.0a2/src/equimo/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.4.0-alpha.2"

{equimo-0.3.5 → equimo-0.4.0a2}/src/equimo/layers/attention.py

@@ -1,4 +1,5 @@
-from typing import Callable, List, Optional, Tuple
+from typing import Callable, List, Literal, Optional, Sequence, Tuple
+
 import equinox as eqx
 import jax
 import jax.numpy as jnp
@@ -6,6 +7,7 @@ import jax.random as jr
 from einops import rearrange, reduce
 from jaxtyping import Array, Float, PRNGKeyArray
 
+from equimo.layers.convolution import ConvBlock, SingleConvBlock, MBConv
 from equimo.layers.dropout import DropPathAdd
 from equimo.layers.ffn import Mlp
 from equimo.layers.mamba import Mamba2Mixer
@@ -1565,6 +1567,487 @@ class LinearAngularAttention(eqx.Module):
         return x
 
 
+class RFAttention(eqx.Module):
+    """Attention with Tensor Reduction by Summation[1].
+
+    A ReLU Linear Attention mechanism replacing matmuls with global
+    summation and element-wise multiplications.
+
+    Attributes:
+        dim: Total dimension of the input/output
+        num_heads: Number of attention heads
+        head_dim: Dimension of each attention head (dim // num_heads)
+
+    References:
+        [1]. Yang, J., An, L., & Park, S. I. (2024). ReduceFormer: Attention
+            with Tensor Reduction by Summation (No. arXiv:2406.07488). arXiv.
+            https://doi.org/10.48550/arXiv.2406.07488
+    """
+
+    total_dim: int = eqx.field(static=True)
+    kernel_func: Callable = eqx.field(static=True)
+    eps: float = eqx.field(static=True)
+
+    qkv: eqx.nn.Conv2d
+    aggreg: list[eqx.nn.Conv2d]
+    proj: SingleConvBlock
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        *,
+        key: PRNGKeyArray,
+        num_heads: int | None = None,
+        head_dim: int = 8,
+        heads_ratio: float = 1.0,
+        scales: Sequence[int] = (5,),
+        use_bias: bool = False,
+        kernel_func: Callable = jax.nn.relu,
+        # TODO: Benchmark against LN, RMSN, NsLN
+        norm_layer: eqx.Module = eqx.nn.GroupNorm,
+        norm_kwargs: dict = {},
+        eps: float = 1e-15,
+        **kwargs,
+    ):
+        key_qkv, key_aggreg, key_proj = jr.split(key, 3)
+
+        self.kernel_func = kernel_func
+        self.eps = eps
+        num_heads = num_heads or int(in_channels // head_dim * heads_ratio)
+        total_dim = num_heads * head_dim
+        self.total_dim = total_dim * (1 + len(scales))
+
+        self.qkv = eqx.nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=3 * total_dim,
+            kernel_size=1,
+            padding="SAME",
+            use_bias=use_bias,
+            key=key_qkv,
+        )
+        self.aggreg = [
+            eqx.nn.Conv2d(
+                in_channels=3 * total_dim,
+                out_channels=3 * total_dim,
+                kernel_size=scale,
+                padding="SAME",
+                groups=3 * total_dim,
+                key=key_aggreg,
+                use_bias=use_bias,
+            )
+            for scale in scales
+        ]
+        # TODO: test different normalizations
+        self.proj = SingleConvBlock(
+            in_channels=self.total_dim,
+            out_channels=out_channels,
+            kernel_size=1,
+            use_bias=use_bias,
+            norm_layer=norm_layer,
+            norm_kwargs=norm_kwargs,
+            key=key_proj,
+        )
+
+    def __call__(
+        self,
+        x: Float[Array, "seqlen height width"],
+        key: PRNGKeyArray,
+        inference: Optional[bool] = None,
+    ) -> Float[Array, "seqlen height width"]:
+        qkv_base = self.qkv(x)
+
+        aggregated_qkvs = [op(qkv_base) for op in self.aggreg]
+        all_qkvs = [qkv_base] + aggregated_qkvs
+
+        rearranged_qkvs = [
+            rearrange(qkv, "(n d) h w -> n d h w", n=3) for qkv in all_qkvs
+        ]
+        multiscale_qkv = jnp.concatenate(rearranged_qkvs, axis=1)
+
+        q, k, v = multiscale_qkv
+
+        q = self.kernel_func(q)
+        k = self.kernel_func(k)
+
+        sum_k = jnp.sum(k, axis=(-1, -2), keepdims=True)
+        sum_v = jnp.sum(v, axis=(-1, -2), keepdims=True)
+        sum_kv = jnp.sum(k * sum_v, axis=(-1, -2), keepdims=True)
+        sum_q = jnp.sum(q, axis=0, keepdims=True)
+
+        out = (q * sum_kv) / (sum_q * sum_k + self.eps)
+        out = self.proj(out)
+
+        return out
+
+
+class RFAttentionBlock(eqx.Module):
+    context_module: RFAttention
+    local_module: MBConv
+
+    def __init__(
+        self,
+        in_channels: int,
+        *,
+        key,
+        head_dim: int = 32,
+        heads_ratio: float = 1.0,
+        scales: Sequence[int] = (5,),
+        rfattn_norm_layer: eqx.Module = eqx.nn.GroupNorm,
+        norm_kwargs: dict = {},
+        expand_ratio: float = 4.0,
+        mbconv_norm_layers: tuple = (None, None, eqx.nn.GroupNorm),
+        mbconv_act_layers: tuple = (jax.nn.hard_swish, jax.nn.hard_swish, None),
+        fuse_mbconv: bool = False,
+        **kwargs,
+    ):
+        key_context, key_local = jr.split(key, 2)
+
+        self.context_module = RFAttention(
+            in_channels=in_channels,
+            out_channels=in_channels,
+            head_dim=head_dim,
+            heads_ratio=heads_ratio,
+            scales=scales,
+            norm_layer=rfattn_norm_layer,
+            norm_kwargs=norm_kwargs,
+            key=key_context,
+        )
+        self.local_module = MBConv(
+            in_channels=in_channels,
+            out_channels=in_channels,
+            expand_ratio=expand_ratio,
+            norm_layers=mbconv_norm_layers,
+            act_layers=mbconv_act_layers,
+            use_bias=(True, True, False),
+            fuse=fuse_mbconv,
+            key=key_local,
+        )
+
+    def __call__(
+        self,
+        x: Float[Array, "dim height width"],
+        key: PRNGKeyArray,
+        inference: Optional[bool] = None,
+    ):
+        key_context, key_local = jr.split(key, 2)
+
+        x += self.context_module(x, inference=inference, key=key_context)
+        x += self.local_module(x, inference=inference, key=key_local)
+
+        return x
+
+
+class ConvAttention(eqx.Module):
+    """Lightweight ConvAttention from LowFormer."""
+
+    num_heads: int = eqx.field(static=True)
+    head_dim: int = eqx.field(static=True)
+    attention_type: Literal["softmax", "sigmoid"] = eqx.field(static=True)
+
+    qkv: eqx.nn.Sequential
+    out_proj: eqx.nn.Conv2d | eqx.nn.Identity
+    upsample: eqx.nn.ConvTranspose2d
+
+    def __init__(
+        self,
+        in_channels: int,
+        *,
+        key: PRNGKeyArray,
+        att_kernel: int = 7,
+        att_stride: int = 4,
+        fuse: bool = True,
+        attention_type: Literal["softmax", "sigmoid"] = "softmax",
+        norm_layer: eqx.Module | None = eqx.nn.GroupNorm,
+        norm_kwargs: dict = {},
+        **kwargs,
+    ):
+        key_qkv1, key_qkv2, key_oproj, key_upsampling = jr.split(key, 4)
+
+        self.attention_type = attention_type
+        self.num_heads = int(max(1, in_channels // 30))
+        self.head_dim = in_channels // self.num_heads
+        total_dim = int(self.head_dim * self.num_heads * 3)
+
+        self.qkv = eqx.nn.Sequential(
+            [
+                SingleConvBlock(
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    kernel_size=att_kernel,
+                    stride=att_stride,
+                    padding=att_kernel // 2,
+                    groups=in_channels,
+                    use_bias=False,
+                    norm_layer=norm_layer,
+                    act_layer=None,
+                    key=key_qkv1,
+                ),
+                eqx.nn.Conv2d(
+                    in_channels=in_channels,
+                    out_channels=total_dim,
+                    kernel_size=1,
+                    stride=1,
+                    padding=0,
+                    use_bias=False,
+                    key=key_qkv2,
+                ),
+            ]
+        )
+
+        if fuse:
+            self.out_proj = eqx.nn.Identity()
+            self.upsample = eqx.nn.ConvTranspose2d(
+                in_channels=self.head_dim * self.num_heads,
+                out_channels=in_channels,
+                kernel_size=3 if att_stride == 1 else (att_stride * 2),
+                stride=att_stride,
+                padding=1 if att_stride == 1 else (att_stride // 2),
+                key=key_upsampling,
+            )
+        else:
+            self.out_proj = eqx.nn.Conv2d(
+                self.head_dim * self.num_heads,
+                in_channels,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                key=key_oproj,
+            )
+            self.upsample = eqx.nn.ConvTranspose2d(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                kernel_size=3 if att_stride == 1 else (att_stride * 2),
+                stride=att_stride,
+                padding=1 if att_stride == 1 else (att_stride // 2),
+                groups=in_channels,
+                key=key_upsampling,
+            )
+
+    def __call__(
+        self,
+        x: Float[Array, "seqlen height width"],
+        key: PRNGKeyArray,
+        inference: Optional[bool] = None,
+    ) -> Float[Array, "seqlen height width"]:
+        q, k, v = rearrange(
+            self.qkv(x),
+            "(n h d) s1 s2 -> n h (s1 s2) d",
+            n=3,
+            h=self.num_heads,
+            d=self.head_dim,
+        )
+
+        _, s, _ = q.shape
+        h = w = int(s**0.5)
+
+        attn_logits = q @ rearrange(k, "h s d -> h d s")
+        attn_logits /= self.head_dim**0.5
+
+        if self.attention_type == "softmax":
+            attn = jax.nn.softmax(attn_logits, axis=-1)
+        elif self.attention_type == "sigmoid":
+            attn = jax.nn.sigmoid(attn_logits)
+
+        v = attn @ v
+
+        out = self.out_proj(rearrange(v, "h (s1 s2) d -> (h d) s1 s2", s1=h, s2=w))
+        out = self.upsample(out)
+
+        return out
+
+
+class ConvAttentionBlock(eqx.Module):
+    prenorm: eqx.Module
+    postnorm: eqx.Module
+    norm: eqx.Module
+    attn: eqx.Module
+    mlp: eqx.Module
+    drop_path1: DropPathAdd
+    drop_path2: DropPathAdd
+
+    def __init__(
+        self,
+        in_channels: int,
+        *,
+        key: PRNGKeyArray,
+        mlp_ratio: float = 4.0,
+        att_stride: int = 1,
+        attention_type: Literal["softmax", "sigmoid"] = "softmax",
+        fuse: bool = True,  # TODO: verify
+        drop_path: float | List[float] = 0.0,
+        act_layer: Callable = jax.nn.gelu,  # TODO: try hardswish
+        norm_layer: eqx.Module = eqx.nn.GroupNorm,
+        norm_max_group: int = 32,
+        post_attention_norm: bool = False,
+        eps: float = 1e-5,
+        **kwargs,
+    ):
+        key_attn, key_mlp = jr.split(key, 2)
+
+        if isinstance(drop_path, list):
+            if len(drop_path) != 2:
+                raise AssertionError(
+                    f"`drop_path` needs to have 2 elements, got {len(drop_path)} ({drop_path})."
+                )
+            dr1, dr2 = drop_path
+            dr1 = float(dr1)
+            dr2 = float(dr2)
+        else:
+            dr1 = dr2 = float(drop_path)
+
+        num_groups = nearest_power_of_2_divisor(in_channels, norm_max_group)
+        self.prenorm = norm_layer(num_groups, in_channels, eps=eps)
+        self.postnorm = (
+            norm_layer(num_groups, in_channels, eps=eps)
+            if post_attention_norm
+            else eqx.nn.Identity()
+        )
+        self.norm = norm_layer(num_groups, in_channels, eps=eps)
+
+        self.attn = ConvAttention(
+            in_channels=in_channels,
+            att_kernel=5 if att_stride > 1 else 3,
+            att_stride=att_stride,
+            fuse=fuse,
+            attention_type=attention_type,
+            key=key_attn,
+        )
+
+        self.mlp = ConvBlock(
+            dim=in_channels,
+            hidden_dim=int(in_channels * mlp_ratio),
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            drop_path=dr1,
+            act_layer=act_layer,
+            key=key_mlp,
+        )
+
+        self.drop_path1 = DropPathAdd(dr1)
+        self.drop_path2 = DropPathAdd(dr2)
+
+    def __call__(
+        self,
+        x: Float[Array, "seqlen dim"],
+        key: PRNGKeyArray,
+        inference: Optional[bool] = None,
+    ) -> Float[Array, "seqlen dim"]:
+        key_attn, key_mlp, key_dr1, key_dr2 = jr.split(key, 4)
+
+        x = self.drop_path1(
+            x,
+            self.postnorm(
+                self.attn(
+                    self.prenorm(x),
+                    inference=inference,
+                    key=key_attn,
+                )
+            ),
+            inference=inference,
+            key=key_dr1,
+        )
+        x = self.drop_path2(
+            x,
+            self.mlp(
+                self.norm(x),
+                inference=inference,
+                key=key_mlp,
+            ),
+            inference=inference,
+            key=key_dr2,
+        )
+
+        return x
+
+
+class LowFormerBlock(eqx.Module):
+    context_module: ConvAttentionBlock
+    local_module: MBConv
+    drop_path1: DropPathAdd
+    drop_path2: DropPathAdd
+
+    def __init__(
+        self,
+        in_channels: int,
+        *,
+        key,
+        mlp_ratio: float = 4.0,
+        att_stride: int = 1,
+        attention_type: Literal["softmax", "sigmoid"] = "softmax",
+        fuse_conv: bool = True,
+        drop_path: float | List[float] = 0.0,
+        act_layer: Callable = jax.nn.hard_swish,
+        norm_layer: eqx.Module = eqx.nn.GroupNorm,
+        expand_ratio: float = 4.0,
+        mbconv_norm_layers: tuple = (None, None, eqx.nn.GroupNorm),
+        mbconv_act_layers: tuple = (jax.nn.hard_swish, jax.nn.hard_swish, None),
+        fuse_mbconv: bool = False,
+        **kwargs,
+    ):
+        key_context, key_local = jr.split(key, 2)
+
+        if isinstance(drop_path, list):
+            if len(drop_path) != 2:
+                raise AssertionError(
+                    f"`drop_path` needs to have 2 elements, got {len(drop_path)} ({drop_path})."
+                )
+            dr1, dr2 = drop_path
+            dr1 = float(dr1)
+            dr2 = float(dr2)
+        else:
+            dr1 = dr2 = float(drop_path)
+
+        self.context_module = ConvAttentionBlock(
+            in_channels=in_channels,
+            mlp_ratio=mlp_ratio,
+            att_stride=att_stride,
+            attention_type=attention_type,
+            fuse=fuse_conv,
+            drop_path=drop_path,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+            key=key_context,
+        )
+        self.local_module = MBConv(
+            in_channels=in_channels,
+            out_channels=in_channels,
+            expand_ratio=expand_ratio,
+            norm_layers=mbconv_norm_layers,
+            act_layers=mbconv_act_layers,
+            use_bias=(True, True, False),
+            fuse=fuse_mbconv,
+            key=key_local,
+        )
+
+        self.drop_path1 = DropPathAdd(dr1)
+        self.drop_path2 = DropPathAdd(dr2)
+
+    def __call__(
+        self,
+        x: Float[Array, "dim height width"],
+        key: PRNGKeyArray,
+        inference: Optional[bool] = None,
+    ):
+        key_context, key_local, key_dr1, key_dr2 = jr.split(key, 4)
+
+        x = self.drop_path1(
+            x,
+            self.context_module(x, inference=inference, key=key_context),
+            inference=inference,
+            key=key_dr1,
+        )
+        x = self.drop_path2(
+            x,
+            self.local_module(x, inference=inference, key=key_local),
+            inference=inference,
+            key=key_dr2,
+        )
+
+        return x
+
+
 def get_attention(module: str | eqx.Module) -> eqx.Module:
     """Get an `eqx.Module` from its common name.