fix: MaskedCausalVisionTransformer build on all supported timm versions

benchmarks/imagenet/vitb16/aim.py

@@ -31,6 +31,7 @@ def __init__(self, batch_size_per_device: int, num_classes: int) -> None:
             qk_norm=False,
             class_token=False,
             no_embed_class=True,
+            global_pool="avg",
         )
         utils.initialize_2d_sine_cosine_positional_embedding(
             pos_embedding=vit.pos_embed, has_class_token=vit.has_class_token

examples/notebooks/pytorch/aim.ipynb

@@ -110,6 +110,7 @@
     "    qk_norm=False,\n",
     "    class_token=False,\n",
     "    no_embed_class=True,\n",
+    "    global_pool=\"avg\",\n",
     ")\n",
     "model = AIM(vit)"
    ]

examples/notebooks/pytorch_lightning/aim.ipynb

@@ -73,6 +73,7 @@
     "            qk_norm=False,\n",
     "            class_token=False,\n",
     "            no_embed_class=True,\n",
+    "            global_pool=\"avg\",\n",
     "        )\n",
     "        utils.initialize_2d_sine_cosine_positional_embedding(\n",
     "            pos_embedding=vit.pos_embed, has_class_token=vit.has_class_token\n",

examples/notebooks/pytorch_lightning_distributed/aim.ipynb

@@ -73,6 +73,7 @@
     "            qk_norm=False,\n",
     "            class_token=False,\n",
     "            no_embed_class=True,\n",
+    "            global_pool=\"avg\",\n",
     "        )\n",
     "        utils.initialize_2d_sine_cosine_positional_embedding(\n",
     "            pos_embedding=vit.pos_embed, has_class_token=vit.has_class_token\n",

examples/pytorch/aim.py

@@ -57,6 +57,7 @@ def forward(self, images):
     qk_norm=False,
     class_token=False,
     no_embed_class=True,
+    global_pool="avg",
 )
 model = AIM(vit)
 

examples/pytorch_lightning/aim.py

@@ -27,6 +27,7 @@ def __init__(self) -> None:
             qk_norm=False,
             class_token=False,
             no_embed_class=True,
+            global_pool="avg",
         )
         utils.initialize_2d_sine_cosine_positional_embedding(
             pos_embedding=vit.pos_embed, has_class_token=vit.has_class_token

examples/pytorch_lightning_distributed/aim.py

@@ -27,6 +27,7 @@ def __init__(self) -> None:
             qk_norm=False,
             class_token=False,
             no_embed_class=True,
+            global_pool="avg",
         )
         utils.initialize_2d_sine_cosine_positional_embedding(
             pos_embedding=vit.pos_embed, has_class_token=vit.has_class_token

lightly/models/modules/masked_causal_vision_transformer.py

@@ -1,3 +1,4 @@
+import inspect
 from typing import Optional
 
 import torch
@@ -120,9 +121,16 @@ class for input parameters.
             *args,
             **kwargs,
         )
+        # Keep only the kwargs that timm's Attention accepts. Block-level kwargs
+        # such as mlp_ratio do not apply to the attention layer.
+        attention_kwargs = {
+            key: value
+            for key, value in kwargs.items()
+            if key in inspect.signature(Attention.__init__).parameters
+        }
         self.attn = MaskedCausalAttention(
             *args,
-            **kwargs,
+            **attention_kwargs,
         )
 
     def forward(

tests/models/modules/test_masked_causal_vision_transformer.py

@@ -0,0 +1,98 @@
+import pytest
+import torch
+
+from lightly.utils import dependency
+
+if not dependency.timm_vit_available():
+    # We do not use pytest.importorskip on module level because it makes mypy unhappy.
+    pytest.skip("TIMM vision transformer is not available", allow_module_level=True)
+
+from lightly.models.modules.masked_causal_vision_transformer import (
+    MaskedCausalAttention,
+    MaskedCausalBlock,
+    MaskedCausalVisionTransformer,
+)
+
+# MaskedCausalAttention supports only fused attention, which needs
+# torch.nn.functional.scaled_dot_product_attention (PyTorch >=2.0).
+skip_without_fused_attention = pytest.mark.skipif(
+    not hasattr(torch.nn.functional, "scaled_dot_product_attention"),
+    reason="MaskedCausalAttention supports only fused attention (PyTorch >=2.0).",
+)
+
+
+class TestMaskedCausalBlock:
+    def test_init__forwards_only_attention_kwargs(self) -> None:
+        # Block-level kwargs such as mlp_ratio must not reach the attention layer.
+        block = MaskedCausalBlock(dim=24, num_heads=3, mlp_ratio=4.0, qkv_bias=True)
+        assert isinstance(block.attn, MaskedCausalAttention)
+        # The mlp_ratio is still applied to the block's mlp, not dropped.
+        fc1 = block.mlp.fc1
+        assert isinstance(fc1, torch.nn.Linear)
+        assert fc1.out_features == 24 * 4
+
+
+class TestMaskedCausalVisionTransformer:
+    def test_init(self) -> None:
+        model = MaskedCausalVisionTransformer(
+            img_size=32,
+            patch_size=16,
+            embed_dim=24,
+            depth=2,
+            num_heads=3,
+            mlp_ratio=4.0,
+        )
+        assert all(
+            isinstance(block.attn, MaskedCausalAttention) for block in model.blocks
+        )
+
+    @skip_without_fused_attention
+    def test_init__aim_config(self) -> None:
+        # The AIM examples build the backbone without a class token and with
+        # average pooling. timm asserts global_pool != "token" when class_token is
+        # False, so both settings are required for the backbone to build.
+        model = MaskedCausalVisionTransformer(
+            img_size=32,
+            patch_size=16,
+            embed_dim=24,
+            depth=2,
+            num_heads=3,
+            class_token=False,
+            no_embed_class=True,
+            global_pool="avg",
+        )
+        assert all(
+            isinstance(block.attn, MaskedCausalAttention) for block in model.blocks
+        )
+        images = torch.rand(2, 3, 32, 32)
+        sequence_length = (32 // 16) ** 2 + model.num_prefix_tokens
+        mask = torch.zeros(2, sequence_length, dtype=torch.bool)
+        mask[:, model.num_prefix_tokens :] = True
+        features = model.forward_features(images, mask=mask)
+        assert features.shape == (2, sequence_length, 24)
+
+    @skip_without_fused_attention
+    def test_forward(self) -> None:
+        model = MaskedCausalVisionTransformer(
+            img_size=32, patch_size=16, embed_dim=24, depth=2, num_heads=3
+        )
+        images = torch.rand(2, 3, 32, 32)
+        features = model.forward_features(images)
+        sequence_length = (32 // 16) ** 2 + model.num_prefix_tokens
+        assert features.shape == (2, sequence_length, 24)
+
+    @skip_without_fused_attention
+    def test_forward__with_mask(self) -> None:
+        model = MaskedCausalVisionTransformer(
+            img_size=32, patch_size=16, embed_dim=24, depth=2, num_heads=3
+        )
+        images = torch.rand(2, 3, 32, 32)
+        sequence_length = (32 // 16) ** 2 + model.num_prefix_tokens
+        mask = torch.zeros(2, sequence_length, dtype=torch.bool)
+        mask[:, model.num_prefix_tokens :] = True
+
+        features = model.forward_features(images, mask=mask)
+        features_no_mask = model.forward_features(images)
+        assert features.shape == (2, sequence_length, 24)
+        # The mask switches the patch tokens to causal attention and changes the output.
+        assert not torch.allclose(features, features_no_mask)