Equimo 0.1.0__tar.gz → 0.1.2__tar.gz

{equimo-0.1.0 → equimo-0.1.2}/Equimo.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: Equimo
-Version: 0.1.0
+Version: 0.1.2
 Summary: Implementation of popular vision models in Jax
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown

{equimo-0.1.0 → equimo-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: Equimo
-Version: 0.1.0
+Version: 0.1.2
 Summary: Implementation of popular vision models in Jax
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown

{equimo-0.1.0 → equimo-0.1.2}/equimo/models/emamodel.py

@@ -32,6 +32,22 @@ class EmaModel:
         self.decay = decay
         self.ema_params = None
 
+    def features(self, *args, **kwargs):
+        """
+        Get the underlying model's features.
+
+        This method relies on the implementation of the `features()` function
+        of the underlying model.
+
+        Args:
+            *args: Positional arguments to pass to the underlying model.
+            **kwargs: Keyword arguments to pass to the underlying model.
+
+        Returns:
+            The output of the underlying model.
+        """
+        return self.model.features(*args, **kwargs)
+
     def __call__(self, *args, **kwargs):
         """
         Call the underlying model.

{equimo-0.1.0 → equimo-0.1.2}/equimo/models/fastervit.py

@@ -431,7 +431,11 @@ class FasterViT(eqx.Module):
 
         num_features = int(dim * 2 ** (len(depths) - 1))
         self.norm = norm_layer(num_features)
-        self.head = eqx.nn.Linear(num_features, num_classes, key=key_head)
+        self.head = (
+            eqx.nn.Linear(num_features, num_classes, key=key_head)
+            if num_classes > 0
+            else eqx.nn.Identity()
+        )
 
     def features(
         self,
@@ -455,6 +459,8 @@ class FasterViT(eqx.Module):
         for blk, key_block in zip(self.blocks, block_subkeys):
             x = blk(x, enable_dropout=enable_dropout, key=key_block)
 
+        x = rearrange(x, "c h w -> (h w) c")
+
         return x
 
     def __call__(
@@ -474,7 +480,6 @@ class FasterViT(eqx.Module):
             Classification logits for each class
         """
         x = self.features(x, enable_dropout, key)
-        x = rearrange(x, "c h w -> (h w) c")
         x = jax.vmap(self.norm)(x)
         x = pool_sd(
             x,

{equimo-0.1.0 → equimo-0.1.2}/equimo/models/mlla.py

@@ -126,17 +126,28 @@ class Mlla(eqx.Module):
             for i, depth in enumerate(depths)
         ]
 
+        self.norm = eqx.nn.LayerNorm(self.num_features)
         self.head = (
-            eqx.nn.Sequential(
-                [
-                    eqx.nn.LayerNorm(self.num_features),
-                    eqx.nn.Linear(self.num_features, num_classes, key=key_head),
-                ]
-            )
+            eqx.nn.Linear(self.num_features, num_classes, key=key_head)
             if num_classes > 0
             else eqx.nn.Identity()
         )
 
+    def features(
+        self,
+        x: Float[Array, "..."],
+        enable_dropout: bool,
+        key: PRNGKeyArray,
+    ) -> Float[Array, "..."]:
+        key_pd, *keys = jr.split(key, 1 + len(self.blocks))
+
+        x = self.patch_embed(x)
+        x = self.pos_drop(x, inference=not enable_dropout, key=key_pd)
+        for i, blk in enumerate(self.blocks):
+            x = blk(x, enable_dropout=enable_dropout, key=keys[i])
+
+        return x
+
     def __call__(
         self,
         x: Float[Array, "..."],
@@ -154,12 +165,8 @@ class Mlla(eqx.Module):
             Output tensor (class logits if num_classes > 0,
             otherwise feature representations)
         """
-        key_pd, *keys = jr.split(key, 1 + len(self.blocks))
-
-        x = self.patch_embed(x)
-        x = self.pos_drop(x, inference=not enable_dropout, key=key_pd)
-        for i, blk in enumerate(self.blocks):
-            x = blk(x, enable_dropout=enable_dropout, key=keys[i])
+        x = self.features(x, enable_dropout=enable_dropout, key=key)
+        x = jax.vmap(self.norm)(x)
         x = reduce(x, "s d -> d", "mean")
         x = self.head(x)
 

{equimo-0.1.0 → equimo-0.1.2}/equimo/models/partialformer.py

@@ -349,13 +349,18 @@ class PartialFormer(eqx.Module):
         ]
 
         self.norm = norm_layer(self.num_features)
-        self.head = eqx.nn.Linear(self.num_features, num_classes, key=key_head)
+        self.head = (
+            eqx.nn.Linear(self.num_features, num_classes, key=key_head)
+            if num_classes > 0
+            else eqx.nn.Identity()
+        )
 
     def features(
         self,
         x: Float[Array, "channels height width"],
         enable_dropout: bool,
         key: PRNGKeyArray,
+        return_qa: bool = False,
     ) -> Float[Array, "seqlen dim"]:
         """Extract features from input image using partial attention.
 
@@ -380,7 +385,9 @@ class PartialFormer(eqx.Module):
                 key=key_block,
             )
 
-        return x, qa
+        if return_qa:
+            return x, qa
+        return x
 
     def __call__(
         self,
@@ -398,7 +405,7 @@ class PartialFormer(eqx.Module):
         Returns:
             Classification logits for each class
         """
-        x, qa = self.features(x, enable_dropout, key)
+        x = self.features(x, enable_dropout, key)
         x = jax.vmap(self.norm)(x)
         x = reduce(x, "n c -> c", "mean")
 

{equimo-0.1.0 → equimo-0.1.2}/equimo/models/shvit.py

@@ -336,17 +336,27 @@ class SHViT(eqx.Module):
             for i, depth in enumerate(depths)
         ]
 
+        self.norm = eqx.nn.LayerNorm(dims[-1])
         self.head = (
-            eqx.nn.Sequential(
-                [
-                    eqx.nn.LayerNorm(dims[-1]),
-                    eqx.nn.Linear(dims[-1], num_classes, key=key_head),
-                ]
-            )
+            eqx.nn.Linear(dims[-1], num_classes, key=key_head)
             if num_classes > 0
             else eqx.nn.Identity()
         )
 
+    def features(
+        self,
+        x: Float[Array, "..."],
+        enable_dropout: bool,
+        key: PRNGKeyArray,
+    ) -> Float[Array, "..."]:
+        keys = jr.split(key, len(self.blocks))
+
+        x = self.patch_embed(x)
+        for i, blk in enumerate(self.blocks):
+            x = blk(x, enable_dropout=enable_dropout, key=keys[i])
+
+        return x
+
     def __call__(
         self,
         x: Float[Array, "..."],
@@ -363,14 +373,9 @@ class SHViT(eqx.Module):
         Returns:
             Classification logits for each class
         """
-        keys = jr.split(key, len(self.blocks))
-
-        x = self.patch_embed(x)
-        for i, blk in enumerate(self.blocks):
-            x = blk(x, enable_dropout=enable_dropout, key=keys[i])
-
+        x = self.features(x, enable_dropout=enable_dropout, key=key)
+        x = jax.vmap(self.norm)(x)
         x = reduce(x, "c h w -> c", "mean")
-
         x = self.head(x)
 
         return x

{equimo-0.1.0 → equimo-0.1.2}/equimo/models/vit.py

@@ -279,7 +279,11 @@ class VisionTransformer(eqx.Module):
         ]
 
         self.norm = norm_layer(dim)
-        self.head = eqx.nn.Linear(dim, num_classes, key=key_head)
+        self.head = (
+            eqx.nn.Linear(dim, num_classes, key=key_head)
+            if num_classes > 0
+            else eqx.nn.Identity()
+        )
 
     def resample_pos_embed(
         self,

{equimo-0.1.0 → equimo-0.1.2}/equimo/models/vssd.py

@@ -144,17 +144,28 @@ class Vssd(eqx.Module):
             for i, depth in enumerate(depths)
         ]
 
+        self.norm = eqx.nn.LayerNorm(self.num_features)
         self.head = (
-            eqx.nn.Sequential(
-                [
-                    eqx.nn.LayerNorm(self.num_features),
-                    eqx.nn.Linear(self.num_features, num_classes, key=key_head),
-                ]
-            )
+            eqx.nn.Linear(self.num_features, num_classes, key=key_head)
             if num_classes > 0
             else eqx.nn.Identity()
         )
 
+    def features(
+        self,
+        x: Float[Array, "..."],
+        enable_dropout: bool,
+        key: PRNGKeyArray,
+    ) -> Float[Array, "..."]:
+        key_pd, *keys = jr.split(key, 1 + len(self.blocks))
+
+        x = self.patch_embed(x)
+        x = self.pos_drop(x, inference=not enable_dropout, key=key_pd)
+        for i, blk in enumerate(self.blocks):
+            x = blk(x, enable_dropout=enable_dropout, key=keys[i])
+
+        return x
+
     def __call__(
         self,
         x: Float[Array, "..."],
@@ -177,12 +188,8 @@ class Vssd(eqx.Module):
         3. Global average pooling
         4. Classification head
         """
-        key_pd, *keys = jr.split(key, 1 + len(self.blocks))
-
-        x = self.patch_embed(x)
-        x = self.pos_drop(x, inference=not enable_dropout, key=key_pd)
-        for i, blk in enumerate(self.blocks):
-            x = blk(x, enable_dropout=enable_dropout, key=keys[i])
+        x = self.features(x, enable_dropout=enable_dropout, key=key)
+        x = jax.vmap(self.norm)(x)
         x = reduce(x, "s d -> d", "mean")
         x = self.head(x)
 

{equimo-0.1.0 → equimo-0.1.2}/pyproject.toml

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