Equimo 0.4.0a2__tar.gz → 0.4.0a4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Equimo
-Version: 0.4.0a2
+Version: 0.4.0a4
 Summary: Implementation of popular vision models in Jax
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers

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

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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Equimo
-Version: 0.4.0a2
+Version: 0.4.0a4
 Summary: Implementation of popular vision models in Jax
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers

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

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

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

@@ -97,7 +97,7 @@ class Attention(eqx.Module):
             attn = jnp.where(mask == 0, jnp.finfo(attn.dtype).min, attn)
 
         attn = jax.nn.softmax(attn, axis=-1)
-        attn = self.attn_drop(attn, inference=inference, key=key1)
+        attn = self.attn_drop(attn, inference=infeoence, key=key1)
 
         x = jnp.einsum("hqk,hkd->hqd", attn, v)
         x = rearrange(x, "h s d -> s (h d)")
@@ -1591,6 +1591,7 @@ class RFAttention(eqx.Module):
     qkv: eqx.nn.Conv2d
     aggreg: list[eqx.nn.Conv2d]
     proj: SingleConvBlock
+    proj_drop: eqx.nn.Dropout
 
     def __init__(
         self,
@@ -1604,6 +1605,7 @@ class RFAttention(eqx.Module):
         scales: Sequence[int] = (5,),
         use_bias: bool = False,
         kernel_func: Callable = jax.nn.relu,
+        proj_drop: float = 0.0,
         # TODO: Benchmark against LN, RMSN, NsLN
         norm_layer: eqx.Module = eqx.nn.GroupNorm,
         norm_kwargs: dict = {},
@@ -1649,12 +1651,14 @@ class RFAttention(eqx.Module):
             key=key_proj,
         )
 
+        self.proj_drop = eqx.nn.Dropout(proj_drop)
+
     def __call__(
         self,
-        x: Float[Array, "seqlen height width"],
+        x: Float[Array, "dim height width"],
         key: PRNGKeyArray,
         inference: Optional[bool] = None,
-    ) -> Float[Array, "seqlen height width"]:
+    ) -> Float[Array, "dim height width"]:
         qkv_base = self.qkv(x)
 
         aggregated_qkvs = [op(qkv_base) for op in self.aggreg]
@@ -1671,12 +1675,13 @@ class RFAttention(eqx.Module):
         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_v = jnp.sum(v * sum_k, 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)
+        out = self.proj_drop(out, inference=inference, key=key)
 
         return out
 
@@ -1684,6 +1689,8 @@ class RFAttention(eqx.Module):
 class RFAttentionBlock(eqx.Module):
     context_module: RFAttention
     local_module: MBConv
+    drop_path1: DropPathAdd
+    drop_path2: DropPathAdd
 
     def __init__(
         self,
@@ -1699,10 +1706,25 @@ class RFAttentionBlock(eqx.Module):
         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,
+        context_drop: float = 0.0,
+        local_drop: float = 0.0,
+        drop_path: float | List[float] = 0.0,
+        residual_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 = RFAttention(
             in_channels=in_channels,
             out_channels=in_channels,
@@ -1711,6 +1733,7 @@ class RFAttentionBlock(eqx.Module):
             scales=scales,
             norm_layer=rfattn_norm_layer,
             norm_kwargs=norm_kwargs,
+            proj_drop=context_drop,
             key=key_context,
         )
         self.local_module = MBConv(
@@ -1721,19 +1744,36 @@ class RFAttentionBlock(eqx.Module):
             act_layers=mbconv_act_layers,
             use_bias=(True, True, False),
             fuse=fuse_mbconv,
+            dropout=local_drop,
+            drop_path=dr1,
+            residual=False,
             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 = jr.split(key, 2)
+        key_context, key_local, key_dr1, key_dr2 = jr.split(key, 4)
 
-        x += self.context_module(x, inference=inference, key=key_context)
-        x += self.local_module(x, inference=inference, key=key_local)
+        # TODO: some prenorm?
+        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
 

{equimo-0.4.0a2 → equimo-0.4.0a4}/src/equimo/layers/convolution.py

@@ -663,11 +663,14 @@ class MBConv(eqx.Module):
     """
 
     fused: bool = eqx.field(static=True)
+    residual: bool = eqx.field(static=True)
 
     inverted_conv: SingleConvBlock | None
     depth_conv: SingleConvBlock | None
     spatial_conv: SingleConvBlock | None
     point_conv: SingleConvBlock
+    dropout: eqx.nn.Dropout
+    drop_path: DropPathAdd
 
     def __init__(
         self,
@@ -688,6 +691,9 @@ class MBConv(eqx.Module):
         fuse_threshold: int = 256,
         fuse_group: bool = False,
         fused_conv_groups: int = 1,
+        dropout: float = 0.0,
+        drop_path: float = 0.0,
+        residual: bool = False,
         **kwargs,
     ):
         key_inverted, key_depth, key_point = jr.split(key, 3)
@@ -711,6 +717,9 @@ class MBConv(eqx.Module):
                 f"`act_layers` should be a Tuple of length 3, got: {len(act_layers)}"
             )
 
+        # Ensure shapes are the same between input and output
+        self.residual = residual and (stride == 1) and (in_channels == out_channels)
+
         mid_channels = (
             mid_channels
             if mid_channels is not None
@@ -779,25 +788,38 @@ class MBConv(eqx.Module):
             key=key_point,
         )
 
+        self.dropout = eqx.nn.Dropout(dropout)
+        self.drop_path = DropPathAdd(drop_path)
+
     def __call__(
         self,
         x: Float[Array, "channels height width"],
         key: PRNGKeyArray,
         inference: Optional[bool] = None,
     ):
+        key_dropout, key_droppath = jr.split(key, 2)
         if self.fused:
-            x = self.spatial_conv(x)
+            out = self.spatial_conv(x)
         else:
-            x = self.inverted_conv(x)
-            x = self.depth_conv(x)
-        x = self.point_conv(x)
+            out = self.inverted_conv(x)
+            out = self.depth_conv(out)
+        out = self.point_conv(out)
 
-        return x
+        out = self.dropout(out, inference=inference, key=key_dropout)
+
+        if self.residual:
+            out = self.drop_path(x, out, inference=inference, key=key_droppath)
+
+        return out
 
 
 class DSConv(eqx.Module):
+    residual: bool = eqx.field(static=True)
+
     depth_conv: SingleConvBlock
     point_conv: SingleConvBlock
+    dropout: eqx.nn.Dropout
+    drop_path: DropPathAdd
 
     def __init__(
         self,
@@ -812,6 +834,9 @@ class DSConv(eqx.Module):
         | eqx.Module
         | None = eqx.nn.GroupNorm,
         act_layers: Tuple[Callable | None, ...] | Callable | None = jax.nn.relu6,
+        residual: bool = False,
+        dropout: float = 0.0,
+        drop_path: float = 0.0,
         **kwargs,
     ):
         key_depth, key_point = jr.split(key, 2)
@@ -835,6 +860,9 @@ class DSConv(eqx.Module):
                 f"`act_layers` should be a Tuple of length 2, got: {len(act_layers)}"
             )
 
+        # Ensure shapes are the same between input and output
+        self.residual = residual and (stride == 1) and (in_channels == out_channels)
+
         self.depth_conv = SingleConvBlock(
             in_channels=in_channels,
             out_channels=in_channels,
@@ -859,13 +887,23 @@ class DSConv(eqx.Module):
             key=key_point,
         )
 
+        self.dropout = eqx.nn.Dropout(dropout)
+        self.drop_path = DropPathAdd(drop_path)
+
     def __call__(
         self,
         x: Float[Array, "channels height width"],
         key: PRNGKeyArray,
         inference: Optional[bool] = None,
     ):
-        x = self.depth_conv(x)
-        x = self.point_conv(x)
+        key_dropout, key_droppath = jr.split(key, 2)
+
+        out = self.depth_conv(x)
+        out = self.point_conv(out)
+
+        out = self.dropout(out, inference=inference, key=key_dropout)
+
+        if self.residual:
+            out = self.drop_path(x, out, inference=inference, key=key_droppath)
 
-        return x
+        return out

{equimo-0.4.0a2 → equimo-0.4.0a4}/src/equimo/layers/dropout.py

@@ -124,5 +124,5 @@ class DropPathAdd(eqx.Module, strict=True):
             )
         else:
             q = 1 - lax.stop_gradient(self.p)
-            add = jrandom.bernouilli(key, q)
+            add = jrandom.bernoulli(key, q)
             return jax.lax.cond(add, lambda x, y: x + y, lambda x, y: x, x1, x2)

{equimo-0.4.0a2 → equimo-0.4.0a4}/src/equimo/models/reduceformer.py

@@ -2,6 +2,7 @@ from typing import Callable, Literal, Optional, Tuple
 
 import equinox as eqx
 import jax
+import jax.numpy as jnp
 import jax.random as jr
 from einops import reduce
 from jaxtyping import Array, Float, PRNGKeyArray
@@ -13,7 +14,6 @@ from equimo.layers.norm import get_norm
 
 
 class BlockChunk(eqx.Module):
-    residuals: list[bool] = eqx.field(static=True)
     blocks: list[DSConv | MBConv | RFAttentionBlock]
 
     def __init__(
@@ -33,6 +33,8 @@ class BlockChunk(eqx.Module):
         act_layer: Callable = jax.nn.hard_swish,
         fewer_norm: bool = False,
         fuse_mbconv: bool = False,
+        dropout: float = 0.0,
+        drop_path: list[float] | float = 0.0,
         **kwargs,
     ):
         key, *block_subkeys = jr.split(key, depth + 1)
@@ -41,8 +43,10 @@ class BlockChunk(eqx.Module):
             k for k, v in kwargs.items() if isinstance(v, list) and len(v) == depth
         ]
 
+        if isinstance(drop_path, float):
+            drop_path = [drop_path] * depth
+
         blocks = []
-        residuals = []
 
         # TODO: simplify logic
         match block_type:
@@ -74,6 +78,9 @@ class BlockChunk(eqx.Module):
                             stride=stride if i == 0 else 1,
                             use_bias=use_bias,
                             norm_layers=norm_layer,
+                            dropout=dropout,
+                            drop_path=drop_path[i],
+                            residual=True,
                             act_layers=(act_layer, None)
                             if block == DSConv
                             else (act_layer, act_layer, None),
@@ -81,9 +88,6 @@ class BlockChunk(eqx.Module):
                             key=block_subkeys[i],
                         )
                     )
-                    residuals.append(
-                        (in_channels == out_channels and stride == 1) or i > 0
-                    )
 
             case "attention":
                 blocks.append(
@@ -95,6 +99,7 @@ class BlockChunk(eqx.Module):
                         norm_layers=(None, None, norm_layer),
                         act_layers=(act_layer, act_layer, None),
                         use_bias=(True, True, False),
+                        dropout=dropout,
                         fuse=fuse_mbconv,
                         key=key,
                     )
@@ -111,13 +116,14 @@ class BlockChunk(eqx.Module):
                             mbconv_norm_layers=(None, None, norm_layer),
                             mbconv_act_layers=(act_layer, act_layer, None),
                             fuse_mbconv=fuse_mbconv,
+                            context_drop=dropout,
+                            local_drop=dropout,
+                            drop_path=drop_path[i],
                             key=block_subkeys[i],
                         )
                     )
-                residuals.append(False)
 
         self.blocks = blocks
-        self.residuals = residuals
 
     def __call__(
         self,
@@ -130,9 +136,8 @@ class BlockChunk(eqx.Module):
         keys = jr.split(key, len(self.blocks))
 
         # TODO: Dropout and Stochastic Path Add
-        for blk, residual, key_block in zip(self.blocks, self.residuals, keys):
-            res = blk(x, inference=inference, key=key_block, **kwargs)
-            x = x + res if residual else res
+        for blk, key_block in zip(self.blocks, keys):
+            x = blk(x, inference=inference, key=key_block, **kwargs)
 
         return x
 
@@ -156,6 +161,9 @@ class ReduceFormer(eqx.Module):
         act_layer: Callable | str = jax.nn.hard_swish,
         fuse_mbconv: bool = False,
         num_classes: int | None = 1000,
+        dropout: float = 0.0,
+        drop_path_rate: float = 0.0,
+        drop_path_uniform: bool = False,
         **kwargs,
     ):
         if not len(widths) == len(depths) == len(block_types):
@@ -165,6 +173,13 @@ class ReduceFormer(eqx.Module):
 
         key_stem, key_head, *key_blocks = jr.split(key, 3 + len(depths))
 
+        depth = sum(depths)
+
+        if drop_path_uniform:
+            dpr = [drop_path_rate] * depth
+        else:
+            dpr = list(jnp.linspace(0.0, drop_path_rate, depth))
+
         act_layer = get_act(act_layer)
         norm_layer = get_norm(norm_layer)
 
@@ -195,6 +210,8 @@ class ReduceFormer(eqx.Module):
                     expand_ratio=1.0,
                     norm_layer=norm_layer,
                     act_layer=act_layer,
+                    dropout=dropout,
+                    drop_path=0.0,
                     key=key_block_stem,
                 ),
             ]
@@ -211,6 +228,8 @@ class ReduceFormer(eqx.Module):
                 norm_layer=norm_layer,
                 act_layer=act_layer,
                 fuse_mbconv=fuse_mbconv,
+                dropout=dropout,
+                drop_path=dpr[sum(depths[:i]) : sum(depths[: i + 1])],
                 key=key_block,
             )
             for i, (depth, block_type, key_block) in enumerate(

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

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