Repository: kornia/kornia

Issue:
Add ability to load official ViT JAX weights directly from Google Bucket
## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->

Currently there is `kornia.contrib.vit` but there are no pretrained weights for them. There are many official JAX checkpoints from Google, and the author also recommends some checkpoints for each model variant.

https://github.com/google-research/vision_transformer

Having the ability to directly load those weights would be convenient. Since the weights are in `.npy` format, we only need numpy to read it. Numpy can be an optional dependency (which means kornia will be able to load JAX weights only if numpy is installed).

Example

```python
from kornia.contrib.vit import VisionTransformer

url = "https://storage.googleapis.com/vit_models/augreg/B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0.npz"

model = VisionTransformer()  # default is ViT-B/16
model.from_jax_weights(url)

# we may also have
from kornia.contrib.vit import vit_b_16

model = vit_b_16(pretrained=True)
```

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->

timm also has a similar function

https://github.com/huggingface/pytorch-image-models/blob/394e8145551191ae60f672556936314a20232a35/timm/models/vision_transformer.py#L723-L727


Discussion:

Resolving patch:
diff --git a/kornia/contrib/vit.py b/kornia/contrib/vit.py
--- a/kornia/contrib/vit.py
+++ b/kornia/contrib/vit.py
@@ -6,12 +6,15 @@
 
 Added some tricks from: `https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py`
 """
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from __future__ import annotations
+
+from typing import Any, Callable
 
 import torch
 from torch import nn
 
-from kornia.core import Module, Tensor
+from kornia.core import Module, Tensor, concatenate
+from kornia.core.check import KORNIA_CHECK
 
 __all__ = ["VisionTransformer"]
 
@@ -21,7 +24,7 @@ def __init__(self, fn: Callable[..., Tensor]) -> None:
         super().__init__()
         self.fn = fn
 
-    def forward(self, x: Tensor, **kwargs: Dict[str, Any]) -> Tensor:
+    def forward(self, x: Tensor, **kwargs: Any) -> Tensor:
         res = x
         x = self.fn(x, **kwargs)
         x += res
@@ -55,12 +58,12 @@ def __init__(self, emb_size: int, num_heads: int, att_drop: float, proj_drop: fl
             )
 
         # fuse the queries, keys and values in one matrix
-        self.qkv = nn.Linear(emb_size, emb_size * 3, bias=False)
+        self.qkv = nn.Linear(emb_size, emb_size * 3)
         self.att_drop = nn.Dropout(att_drop)
         self.projection = nn.Linear(emb_size, emb_size)
         self.projection_drop = nn.Dropout(proj_drop)  # added timm trick
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: Tensor) -> Tensor:
         B, N, C = x.shape
         # split keys, queries and values in num_heads
         # NOTE: the line below differs from timm
@@ -86,15 +89,15 @@ def __init__(self, embed_dim: int, num_heads: int, dropout_rate: float, dropout_
         super().__init__(
             ResidualAdd(
                 nn.Sequential(
-                    nn.LayerNorm(embed_dim),
+                    nn.LayerNorm(embed_dim, 1e-6),
                     MultiHeadAttention(embed_dim, num_heads, dropout_attn, dropout_rate),
                     nn.Dropout(dropout_rate),
                 )
             ),
             ResidualAdd(
                 nn.Sequential(
-                    nn.LayerNorm(embed_dim),
-                    FeedForward(embed_dim, embed_dim, embed_dim, dropout_rate=dropout_rate),
+                    nn.LayerNorm(embed_dim, 1e-6),
+                    FeedForward(embed_dim, embed_dim * 4, embed_dim, dropout_rate=dropout_rate),
                     nn.Dropout(dropout_rate),
                 )
             ),
@@ -114,9 +117,9 @@ def __init__(
         self.blocks = nn.Sequential(
             *(TransformerEncoderBlock(embed_dim, num_heads, dropout_rate, dropout_attn) for _ in range(depth))
         )
-        self.results: List[torch.Tensor] = []
+        self.results: list[Tensor] = []
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: Tensor) -> Tensor:
         self.results = []
         out = x
         for m in self.blocks.children():
@@ -134,7 +137,7 @@ def __init__(
         out_channels: int = 768,
         patch_size: int = 16,
         image_size: int = 224,
-        backbone: Optional[Module] = None,
+        backbone: Module | None = None,
     ) -> None:
         super().__init__()
         self.in_channels = in_channels
@@ -152,17 +155,17 @@ def __init__(
         self.cls_token = nn.Parameter(torch.randn(1, 1, out_channels))
         self.positions = nn.Parameter(torch.randn(feat_size + 1, out_channels))
 
-    def _compute_feats_dims(self, image_size: Tuple[int, int, int]) -> Tuple[int, int]:
+    def _compute_feats_dims(self, image_size: tuple[int, int, int]) -> tuple[int, int]:
         out = self.backbone(torch.zeros(1, *image_size)).detach()
         return out.shape[-3], out.shape[-2] * out.shape[-1]
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: Tensor) -> Tensor:
         x = self.backbone(x)
         B, N, _, _ = x.shape
         x = x.view(B, N, -1).permute(0, 2, 1)  # BxNxE
         cls_tokens = self.cls_token.repeat(B, 1, 1)  # Bx1xE
         # prepend the cls token to the input
-        x = torch.cat([cls_tokens, x], dim=1)  # Bx(N+1)xE
+        x = concatenate([cls_tokens, x], dim=1)  # Bx(N+1)xE
         # add position embedding
         x += self.positions
         return x
@@ -206,7 +209,7 @@ def __init__(
         num_heads: int = 12,
         dropout_rate: float = 0.0,
         dropout_attn: float = 0.0,
-        backbone: Optional[Module] = None,
+        backbone: Module | None = None,
     ) -> None:
         super().__init__()
         self.image_size = image_size
@@ -217,14 +220,15 @@ def __init__(
         self.patch_embedding = PatchEmbedding(in_channels, embed_dim, patch_size, image_size, backbone)
         hidden_dim = self.patch_embedding.out_channels
         self.encoder = TransformerEncoder(hidden_dim, depth, num_heads, dropout_rate, dropout_attn)
+        self.norm = nn.LayerNorm(hidden_dim, 1e-6)
 
     @property
-    def encoder_results(self) -> List[Tensor]:
+    def encoder_results(self) -> list[Tensor]:
         return self.encoder.results
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if not isinstance(x, torch.Tensor):
-            raise TypeError(f"Input x type is not a torch.Tensor. Got: {type(x)}")
+    def forward(self, x: Tensor) -> Tensor:
+        if not isinstance(x, Tensor):
+            raise TypeError(f"Input x type is not a Tensor. Got: {type(x)}")
 
         if self.image_size not in (*x.shape[-2:],) and x.shape[-3] != self.in_channels:
             raise ValueError(
@@ -233,4 +237,57 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         out = self.patch_embedding(x)
         out = self.encoder(out)
+        out = self.norm(out)
         return out
+
+    @staticmethod
+    def from_config(variant: str, pretrained: bool = False, **kwargs: Any) -> VisionTransformer:
+        """Build ViT model based on the given config string. The format is ``vit_{size}/{patch_size}``.
+        E.g. ``vit_b/16`` means ViT-Base, patch size 16x16. If ``pretrained=True``, AugReg weights are loaded.
+        The weights are hosted on HuggingFace's model hub: https://huggingface.co/kornia.
+
+        .. note::
+            The available weights are: ``vit_l/16``, ``vit_b/16``, ``vit_s/16``, ``vit_ti/16``,
+            ``vit_b/32``, ``vit_s/32``.
+
+        Args:
+            variant: ViT model variant e.g. ``vit_b/16``.
+            pretrained: whether to load pre-trained AugReg weights.
+            kwargs: other keyword arguments that will be passed to :func:`kornia.contrib.vit.VisionTransformer`.
+        Returns:
+            The respective ViT model
+
+        Example:
+            >>> from kornia.contrib import VisionTransformer
+            >>> vit_model = VisionTransformer.from_config("vit_b/16", pretrained=True)
+        """
+        model_type, patch_size_str = variant.split("/")
+        patch_size = int(patch_size_str)
+
+        model_config = {
+            "vit_ti": {"embed_dim": 192, "depth": 12, "num_heads": 3},
+            "vit_s": {"embed_dim": 384, "depth": 12, "num_heads": 6},
+            "vit_b": {"embed_dim": 768, "depth": 12, "num_heads": 12},
+            "vit_l": {"embed_dim": 1024, "depth": 24, "num_heads": 16},
+            "vit_h": {"embed_dim": 1280, "depth": 32, "num_heads": 16},
+        }[model_type]
+        kwargs.update(model_config, patch_size=patch_size)
+
+        model = VisionTransformer(**kwargs)
+
+        if pretrained:
+            url = _get_weight_url(variant)
+            state_dict = torch.hub.load_state_dict_from_url(url)
+            model.load_state_dict(state_dict)
+
+        return model
+
+
+_AVAILABLE_WEIGHTS = ["vit_l/16", "vit_b/16", "vit_s/16", "vit_ti/16", "vit_b/32", "vit_s/32"]
+
+
+def _get_weight_url(variant: str) -> str:
+    """Return the URL of the model weights."""
+    KORNIA_CHECK(variant in _AVAILABLE_WEIGHTS, f"Variant {variant} does not have pre-trained checkpoint")
+    model_type, patch_size = variant.split("/")
+    return f"https://huggingface.co/kornia/{model_type}{patch_size}_augreg_i21k_r224/resolve/main/{model_type}-{patch_size}.pth"


Test patch:
diff --git a/tests/contrib/test_contrib.py b/tests/contrib/test_contrib.py
--- a/tests/contrib/test_contrib.py
+++ b/tests/contrib/test_contrib.py
@@ -76,6 +76,16 @@ def backbone_mock(x):
         out = vit(img)
         assert out.shape == (1, 197, 128)
 
+    @pytest.mark.parametrize("variant", ["vit_ti/16", "vit_b/16"])
+    def test_from_config(self, variant):
+        model = kornia.contrib.VisionTransformer.from_config(variant)
+        assert isinstance(model, kornia.contrib.VisionTransformer)
+
+    @pytest.mark.parametrize("variant", ["vit_ti/16"])
+    def test_from_config_pretrained(self, variant):
+        model = kornia.contrib.VisionTransformer.from_config(variant, pretrained=True)
+        assert isinstance(model, kornia.contrib.VisionTransformer)
+
 
 class TestMobileViT(BaseTester):
     @pytest.mark.parametrize("B", [1, 2])