Equimo 0.1.2__tar.gz → 0.1.3a2__tar.gz

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

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

{equimo-0.1.2 → equimo-0.1.3a2}/PKG-INFO

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

{equimo-0.1.2 → equimo-0.1.3a2}/equimo/layers/ffn.py

@@ -9,6 +9,133 @@ from jaxtyping import Array, Float, PRNGKeyArray
 from equimo.layers.dropout import Dropout
 
 
+class WeightNormLinear(eqx.Module):
+    """Linear layer with weight normalization.
+
+    Implements weight normalization as described in "Weight Normalization: A Simple
+    Reparameterization to Accelerate Training of Deep Neural Networks". The weights
+    are parameterized as w = g * (v/||v||), where g is a scalar and v is a vector.
+
+    Attributes:
+        weight_v: The weight vector v to be normalized
+        weight_g: The scalar multiplier g for controlling the output scale
+
+    References:
+        [1] https://arxiv.org/abs/1602.07868
+    """
+
+    weight_v: jnp.ndarray
+    weight_g: jnp.ndarray
+
+    def __init__(self, in_features: int, out_features: int, key: PRNGKeyArray):
+        self.weight_v = eqx.nn.Linear(
+            in_features, out_features, use_bias=False, key=key
+        ).weight
+        self.weight_g = jnp.ones((out_features, 1))
+
+    def __call__(self, x):
+        """Apply weight normalized linear transformation.
+
+        Args:
+            x: Input tensor
+
+        Returns:
+            Transformed tensor using normalized weights
+        """
+        v_norm = jnp.linalg.norm(self.weight_v, ord=2, axis=1, keepdims=True)
+        normalized_v = self.weight_v / v_norm
+        weight = self.weight_g * normalized_v
+
+        out = x @ weight.T
+        return out
+
+
+class DINOHead(eqx.Module):
+    """Multi-layer Perceptron (MLP) head used to train DINOv2 models.
+
+    Implements the projection head architecture from DINOv2 self-supervised learning.
+    Features multiple fully connected layers with activation, followed by L2
+    normalization and a weight-normalized final layer.
+
+    The architecture follows:
+    input -> fc1 -> act -> fc2 -> act -> fc3 -> act -> L2norm -> weight_norm_linear
+
+    Attributes:
+        fc1: First linear layer projecting to hidden dimension
+        fc2: Second linear layer maintaining hidden dimension
+        fc3: Third linear layer projecting to bottleneck dimension
+        last: Final weight-normalized linear layer
+        act_layer: Activation function used between layers
+
+    References:
+        [1] https://arxiv.org/abs/2304.07193
+    """
+
+    fc1: eqx.nn.Linear
+    fc2: eqx.nn.Linear
+    fc3: eqx.nn.Linear
+    last: WeightNormLinear
+    act_layer: Callable
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        *,
+        hidden_dim=2048,
+        bottleneck_dim=256,
+        key: PRNGKeyArray,
+        act_layer: Callable = jax.nn.gelu,
+        **kwargs,
+    ):
+        """Initialize the DINOv2 projection head.
+
+        Args:
+            in_features: Number of input features
+            out_features: Number of output features
+            hidden_dim: Dimension of hidden layers (default: 2048)
+            bottleneck_dim: Dimension of bottleneck layer (default: 256)
+            key: PRNG key for initialization
+            act_layer: Activation function (default: gelu)
+            **kwargs: Additional arguments
+        """
+        key_fc1, key_fc2, key_fc3, key_last = jr.split(key, 3)
+
+        self.act_layer = act_layer
+
+        self.fc1 = eqx.nn.Linear(in_features, hidden_dim, key=key_fc1)
+        self.fc2 = eqx.nn.Linear(hidden_dim, hidden_dim, key=key_fc2)
+        self.fc3 = eqx.nn.Linear(hidden_dim, bottleneck_dim, key=key_fc3)
+        self.last = WeightNormLinear(bottleneck_dim, out_features, key=key_last)
+
+    def __call__(
+        self,
+        x: Float[Array, "seqlen dim"],
+        enable_dropout: bool,
+        key: PRNGKeyArray,
+    ) -> Float[Array, "seqlen dim"]:
+        """Process input through the DINOv2 projection head.
+
+        Args:
+            x: Input feature tensor
+            enable_dropout: Whether to enable dropout (unused in original implementation)
+            key: PRNG key for random operations
+
+        Returns:
+            Projected and normalized features
+        """
+        eps = jnp.where(x.dtype == jnp.float16, 1e-6, 1e-12)
+        x = self.act_layer(jax.vmap(self.fc1)(x))
+        x = self.act_layer(jax.vmap(self.fc2)(x))
+        x = self.act_layer(jax.vmap(self.fc3)(x))
+
+        x = x / (jnp.linalg.norm(x, ord=2, axis=-1, keepdims=True) + eps)
+
+        x = self.last(x)
+
+        return x
+
+
 class Mlp(eqx.Module):
     """Multi-layer Perceptron (MLP) module with dropout.
 

{equimo-0.1.2 → equimo-0.1.3a2}/equimo/models/vit.py

@@ -6,7 +6,7 @@ import jax
 import jax.numpy as jnp
 import jax.random as jr
 from einops import rearrange
-from jaxtyping import Array, Float, PRNGKeyArray
+from jaxtyping import Array, Float, Int, PRNGKeyArray
 
 from equimo.layers.attention import Attention, AttentionBlock
 from equimo.layers.dropout import Dropout
@@ -148,12 +148,14 @@ class VisionTransformer(eqx.Module):
     pos_embed: jnp.ndarray
     cls_token: jnp.ndarray | None
     reg_tokens: jnp.ndarray | None
+    mask_token: jnp.ndarray | None
     blocks: List[eqx.Module]
     pos_drop: Dropout
     norm: eqx.Module
     head: eqx.Module
 
     dim: int = eqx.field(static=True)
+    embed_size: int = eqx.field(static=True)
     num_patches: int = eqx.field(static=True)
     global_pool: str = eqx.field(static=True)
     num_reg_tokens: int = eqx.field(static=True)
@@ -175,6 +177,7 @@ class VisionTransformer(eqx.Module):
         depths: List[int],
         *,
         key: PRNGKeyArray,
+        use_mask_token: bool = False,
         repeat: int = 1,
         dynamic_img_size: bool = False,
         dynamic_img_pad: bool = False,
@@ -217,6 +220,7 @@ class VisionTransformer(eqx.Module):
         self.no_embed_class = no_embed_class
         self.pos_embed_reg_tokens = pos_embed_reg_tokens
         self.global_pool = global_pool
+        self.embed_size = img_size // patch_size
 
         self.patch_embed = PatchEmbedding(
             in_channels,
@@ -232,6 +236,8 @@ class VisionTransformer(eqx.Module):
         self.cls_token = jr.normal(key_cls, (1, dim))
         self.reg_tokens = jr.normal(key_reg, (reg_tokens, dim))
 
+        self.mask_token = jnp.zeros((1, dim)) if use_mask_token else None
+
         if no_embed_class:
             self.embed_len = self.num_patches
         elif self.pos_embed_reg_tokens:
@@ -390,6 +396,7 @@ class VisionTransformer(eqx.Module):
         x: Float[Array, "channels height width"],
         enable_dropout: bool,
         key: PRNGKeyArray,
+        mask: Optional[Int[Array, "embed_h embed_w"]] = None,
     ) -> Float[Array, "seqlen dim"]:
         """Extract features from input image.
 
@@ -397,6 +404,7 @@ class VisionTransformer(eqx.Module):
             x: Input image tensor
             enable_dropout: Whether to enable dropout during inference
             key: PRNG key for random operations
+            mask: optional binary mask of the size of the input after patch embedding
 
         Returns:
             Processed feature tensor
@@ -405,6 +413,16 @@ class VisionTransformer(eqx.Module):
         x = self.patch_embed(x)
         x = self.pos_drop(x, inference=not enable_dropout, key=key_posdrop)
 
+        if mask is not None:
+            assert (
+                self.mask_token is not None
+            ), "To use masked forward, init the model with `use_mask_token=True`."
+            x = jnp.where(
+                rearrange(mask, "h w -> (h w) 1"), x, self.mask_token.astype(x.dtype)
+            )
+        # TODO: Decompose in multiple fns
+        x = self._pos_embed(x, h=self.embed_size, w=self.embed_size)
+
         for blk, key_block in zip(self.blocks, block_subkeys):
             x = blk(x, enable_dropout=enable_dropout, key=key_block)
 
@@ -430,13 +448,7 @@ class VisionTransformer(eqx.Module):
                 - x_norm_patchtokens: Normalized patch tokens
                 - x_prenorm: Pre-normalized features
         """
-        key_posdrop, *block_subkeys = jr.split(key, len(self.blocks) + 1)
-        x = self.patch_embed(x)
-        x = self.pos_drop(x, inference=not enable_dropout, key=key_posdrop)
-
-        for blk, key_block in zip(self.blocks, block_subkeys):
-            x = blk(x, enable_dropout=enable_dropout, key=key_block)
-
+        x = self.features(x, enable_dropout=enable_dropout, key=key)
         x_norm = jax.vmap(self.norm)(x)
 
         return {
@@ -474,3 +486,26 @@ class VisionTransformer(eqx.Module):
         x = self.head(x)
 
         return x
+
+
+#
+# key = jr.PRNGKey(42)
+# img_size = 224
+# patch_size = 14
+#
+# x = jr.normal(key, (3, 224, 224))
+# mask = jr.bernoulli(key, shape=(16, 16)) * 1
+#
+# base_model = VisionTransformer(
+#     img_size=img_size,
+#     in_channels=3,
+#     dim=384,
+#     patch_size=patch_size,
+#     num_heads=[6],
+#     depths=[12],
+#     num_classes=0,
+#     use_mask_token=True,
+#     key=key,
+# )
+# f = base_model.features(x, mask=mask, enable_dropout=True, key=key)
+# self = base_model

{equimo-0.1.2 → equimo-0.1.3a2}/pyproject.toml

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