Cache t_x/t_y for rope mixed, use more like other rope embeds, rename rope_mixed -> mrope in model names

Author: rwightman

timm/layers/pos_embed_sincos.py

@@ -365,9 +365,8 @@ def __init__(
             )
 
     def get_embed(self, shape: Optional[List[int]] = None):
-        if self.bands is not None:
+        if shape is not None and self.bands is not None:
             # rebuild embeddings every call, use if target shape changes
-            assert shape is not None
             return build_rotary_pos_embed(
                 shape,
                 self.bands,
@@ -376,8 +375,10 @@ def get_embed(self, shape: Optional[List[int]] = None):
                 grid_offset=self.grid_offset,
                 grid_indexing=self.grid_indexing,
             )
-        else:
+        elif self.pos_embed_sin is not None and self.pos_embed_cos is not None:
             return self.pos_embed_sin, self.pos_embed_cos
+        else:
+            assert False, "get_embed() requires pre-computed pos embeds or valid shape w/ pre-computed bands"
 
     def forward(self, x):
         # assuming channel-first tensor where spatial dim are >= 2
@@ -456,7 +457,7 @@ def __init__(
             )
 
     def get_embed(self, shape: Optional[List[int]] = None):
-        if self.bands is not None and shape is not None:
+        if shape is not None and self.bands is not None:
             # rebuild embeddings every call, use if target shape changes
             embeds = build_rotary_pos_embed(
                 shape,
@@ -470,7 +471,7 @@ def get_embed(self, shape: Optional[List[int]] = None):
         elif self.pos_embed is not None:
             return self.pos_embed
         else:
-            assert False, "get_embed() requires pre-computed pos_embed or valid shape w/ pre-computed bands"
+            assert False, "get_embed() requires pre-computed pos embed or valid shape w/ pre-computed bands"
 
     def forward(self, x):
         # assuming channel-first tensor where spatial dim are >= 2
@@ -514,31 +515,39 @@ def init_random_2d_freqs(
 
 @torch.fx.wrap
 @register_notrace_function
-def get_mixed_freqs(
-        freqs: torch.Tensor,
+def get_mixed_grid(
         height: int,
         width: int,
         grid_indexing: str = 'ij',
-):
-    """Compute mixed (learnable) frequencies."""
-    # Create position indices
-    device = freqs.device
-    dtype = freqs.dtype
+        device: Optional[torch.device] = None,
+        dtype: torch.dtype = torch.float32,
+) -> Tuple[torch.Tensor, torch.Tensor]:
     x_pos, y_pos = torch.meshgrid(
         torch.arange(height, dtype=dtype, device=device),
         torch.arange(width, dtype=dtype, device=device),
         indexing=grid_indexing,
     )
     t_x = x_pos.flatten()
     t_y = y_pos.flatten()
+    return t_x, t_y
+
+
+def get_mixed_freqs(
+        freqs: torch.Tensor,
+        t_x: torch.Tensor,
+        t_y: torch.Tensor,
+) -> torch.Tensor:
+    """Compute mixed (learnable) frequencies."""
+    # Create position indices
+    dtype = freqs.dtype
+    freqs = freqs.float()
     freqs_x = (t_x.unsqueeze(-1) @ freqs[0].unsqueeze(-2))
     freqs_y = (t_y.unsqueeze(-1) @ freqs[1].unsqueeze(-2))
     combined = freqs_x + freqs_y  # shape: (num_heads, N, dim//4)
     sin_emb = torch.sin(combined).repeat_interleave(2, -1)  # (N, dim//2)
     cos_emb = torch.cos(combined).repeat_interleave(2, -1)  # (N, dim//2)
     rope_embeds = torch.cat([sin_emb, cos_emb], dim=-1)  # (num_heads, H*W, head_dim)
-
-    return rope_embeds
+    return rope_embeds.to(dtype)
 
 
 class RotaryEmbeddingMixed(nn.Module):
@@ -584,6 +593,18 @@ def __init__(
             rotate=True,
         )  # (2, depth, num_heads, head_dim//2)
         self.freqs = nn.Parameter(freqs)
+        if feat_shape is not None:
+            # cache pre-computed grid
+            t_x, t_y = get_mixed_grid(
+                feat_shape[0],
+                feat_shape[1],
+                grid_indexing=grid_indexing,
+                device=self.freqs.device
+            )
+            self.register_buffer('t_x', t_x, persistent=False)
+            self.register_buffer('t_y', t_y, persistent=False)
+        else:
+            self.t_x = self.t_y = None
 
     def get_embed(self, shape: Optional[List[int]] = None) -> torch.Tensor:
         """Generate rotary embeddings for the given spatial shape.
@@ -594,9 +615,19 @@ def get_embed(self, shape: Optional[List[int]] = None) -> torch.Tensor:
         Returns:
             Tensor of shape (depth, H*W, dim) containing concatenated sin/cos embeddings
         """
-        assert shape is not None, "shape must be provided"
-        H, W = shape
-        return get_mixed_freqs(self.freqs, height=H, width=W, grid_indexing=self.grid_indexing)
+        if shape is not None:
+            t_x, t_y = get_mixed_grid(
+                shape[0],
+                shape[1],
+                grid_indexing=self.grid_indexing,
+                device=self.freqs.device
+            )
+        elif self.t_x is not None and self.t_y is not None:
+            t_x, t_y = self.t_x, self.t_y
+        else:
+            assert False, "get_embed() requires pre-computed t_x/t_y or valid shape"
+
+        return get_mixed_freqs(self.freqs, t_x, t_y)
 
     def forward(self, x):
         # assuming channel-first tensor where spatial dim are >= 2

timm/models/eva.py

@@ -708,7 +708,7 @@ def _pos_embed(self, x) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
             rot_pos_embed = self.rope.get_embed(shape=(H, W)) if self.rope is not None else None
         else:
             pos_embed = self.pos_embed
-            rot_pos_embed = self.rope.get_embed(shape=self.patch_embed.grid_size) if self.rope is not None else None
+            rot_pos_embed = self.rope.get_embed() if self.rope is not None else None
 
         to_cat = []
         if self.cls_token is not None:
@@ -1392,19 +1392,19 @@ def _pe_cfg(url: str = '', **kwargs) -> Dict[str, Any]:
         mean=IMAGENET_DEFAULT_MEAN,
         std=IMAGENET_DEFAULT_STD,
     ),
-    'vit_small_patch16_rope_mixed_224.naver_in1k': _cfg(
+    'vit_small_patch16_mrope_224.naver_in1k': _cfg(
         hf_hub_id='naver-ai/rope_mixed_deit_small_patch16_LS',
         hf_hub_filename='pytorch_model.bin',
         mean=IMAGENET_DEFAULT_MEAN,
         std=IMAGENET_DEFAULT_STD,
     ),
-    'vit_base_patch16_rope_mixed_224.naver_in1k': _cfg(
+    'vit_base_patch16_mrope_224.naver_in1k': _cfg(
         hf_hub_id='naver-ai/rope_mixed_deit_base_patch16_LS',
         hf_hub_filename='pytorch_model.bin',
         mean=IMAGENET_DEFAULT_MEAN,
         std=IMAGENET_DEFAULT_STD,
     ),
-    'vit_large_patch16_rope_mixed_224.naver_in1k': _cfg(
+    'vit_large_patch16_mrope_224.naver_in1k': _cfg(
         hf_hub_id='naver-ai/rope_mixed_deit_large_patch16_LS',
         hf_hub_filename='pytorch_model.bin',
         mean=IMAGENET_DEFAULT_MEAN,
@@ -1428,19 +1428,19 @@ def _pe_cfg(url: str = '', **kwargs) -> Dict[str, Any]:
         mean=IMAGENET_DEFAULT_MEAN,
         std=IMAGENET_DEFAULT_STD,
     ),
-    'vit_small_patch16_rope_mixed_ape_224.naver_in1k': _cfg(
+    'vit_small_patch16_mrope_ape_224.naver_in1k': _cfg(
         hf_hub_id='naver-ai/rope_mixed_ape_deit_small_patch16_LS',
         hf_hub_filename='pytorch_model.bin',
         mean=IMAGENET_DEFAULT_MEAN,
         std=IMAGENET_DEFAULT_STD,
     ),
-    'vit_base_patch16_rope_mixed_ape_224.naver_in1k': _cfg(
+    'vit_base_patch16_mrope_ape_224.naver_in1k': _cfg(
         hf_hub_id='naver-ai/rope_mixed_ape_deit_base_patch16_LS',
         hf_hub_filename='pytorch_model.bin',
         mean=IMAGENET_DEFAULT_MEAN,
         std=IMAGENET_DEFAULT_STD,
     ),
-    'vit_large_patch16_rope_mixed_ape_224.naver_in1k': _cfg(
+    'vit_large_patch16_mrope_ape_224.naver_in1k': _cfg(
         hf_hub_id='naver-ai/rope_mixed_ape_deit_large_patch16_LS',
         hf_hub_filename='pytorch_model.bin',
         mean=IMAGENET_DEFAULT_MEAN,
@@ -2023,7 +2023,7 @@ def vit_large_patch16_rope_224(pretrained: bool = False, **kwargs) -> Eva:
 
 
 @register_model
-def vit_small_patch16_rope_mixed_224(pretrained: bool = False, **kwargs) -> Eva:
+def vit_small_patch16_mrope_224(pretrained: bool = False, **kwargs) -> Eva:
     """RoPE-Mixed ViT-S/16 from https://github.com/naver-ai/rope-vit"""
     model_args = dict(
         patch_size=16,
@@ -2042,12 +2042,12 @@ def vit_small_patch16_rope_mixed_224(pretrained: bool = False, **kwargs) -> Eva:
         rope_temperature=10.0,
         rope_mixed_mode=True,
     )
-    model = _create_eva('vit_small_patch16_rope_mixed_224', pretrained=pretrained, **dict(model_args, **kwargs))
+    model = _create_eva('vit_small_patch16_mrope_224', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def vit_base_patch16_rope_mixed_224(pretrained: bool = False, **kwargs) -> Eva:
+def vit_base_patch16_mrope_224(pretrained: bool = False, **kwargs) -> Eva:
     """RoPE-Mixed ViT-B/16 from https://github.com/naver-ai/rope-vit"""
     model_args = dict(
         patch_size=16,
@@ -2066,12 +2066,12 @@ def vit_base_patch16_rope_mixed_224(pretrained: bool = False, **kwargs) -> Eva:
         rope_temperature=10.0,
         rope_mixed_mode=True,
     )
-    model = _create_eva('vit_base_patch16_rope_mixed_224', pretrained=pretrained, **dict(model_args, **kwargs))
+    model = _create_eva('vit_base_patch16_mrope_224', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def vit_large_patch16_rope_mixed_224(pretrained: bool = False, **kwargs) -> Eva:
+def vit_large_patch16_mrope_224(pretrained: bool = False, **kwargs) -> Eva:
     """RoPE-Mixed ViT-L/16 from https://github.com/naver-ai/rope-vit"""
     model_args = dict(
         patch_size=16,
@@ -2090,7 +2090,7 @@ def vit_large_patch16_rope_mixed_224(pretrained: bool = False, **kwargs) -> Eva:
         rope_temperature=10.0,
         rope_mixed_mode=True,
     )
-    model = _create_eva('vit_large_patch16_rope_mixed_224', pretrained=pretrained, **dict(model_args, **kwargs))
+    model = _create_eva('vit_large_patch16_mrope_224', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
@@ -2170,7 +2170,7 @@ def vit_large_patch16_rope_ape_224(pretrained: bool = False, **kwargs) -> Eva:
 
 
 @register_model
-def vit_small_patch16_rope_mixed_ape_224(pretrained: bool = False, **kwargs) -> Eva:
+def vit_small_patch16_mrope_ape_224(pretrained: bool = False, **kwargs) -> Eva:
     """RoPE-Mixed + APE ViT-S/16 from https://github.com/naver-ai/rope-vit"""
     model_args = dict(
         patch_size=16,
@@ -2191,12 +2191,12 @@ def vit_small_patch16_rope_mixed_ape_224(pretrained: bool = False, **kwargs) ->
         rope_mixed_mode=True,
     )
 
-    model = _create_eva('vit_small_patch16_rope_mixed_ape_224', pretrained=pretrained, **dict(model_args, **kwargs))
+    model = _create_eva('vit_small_patch16_mrope_ape_224', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def vit_base_patch16_rope_mixed_ape_224(pretrained: bool = False, **kwargs) -> Eva:
+def vit_base_patch16_mrope_ape_224(pretrained: bool = False, **kwargs) -> Eva:
     """RoPE-Mixed + APE ViT-B/16 from https://github.com/naver-ai/rope-vit"""
     model_args = dict(
         patch_size=16,
@@ -2216,12 +2216,12 @@ def vit_base_patch16_rope_mixed_ape_224(pretrained: bool = False, **kwargs) -> E
         rope_temperature=10.0,
         rope_mixed_mode=True,
     )
-    model = _create_eva('vit_base_patch16_rope_mixed_ape_224', pretrained=pretrained, **dict(model_args, **kwargs))
+    model = _create_eva('vit_base_patch16_mrope_ape_224', pretrained=pretrained, **dict(model_args, **kwargs))
     return model
 
 
 @register_model
-def vit_large_patch16_rope_mixed_ape_224(pretrained: bool = False, **kwargs) -> Eva:
+def vit_large_patch16_mrope_ape_224(pretrained: bool = False, **kwargs) -> Eva:
     """RoPE-Mixed + APE ViT-L/16 from https://github.com/naver-ai/rope-vit"""
     model_args = dict(
         patch_size=16,
@@ -2241,6 +2241,6 @@ def vit_large_patch16_rope_mixed_ape_224(pretrained: bool = False, **kwargs) ->
         rope_temperature=10.0,
         rope_mixed_mode=True,
     )
-    model = _create_eva('vit_large_patch16_rope_mixed_ape_224', pretrained=pretrained, **dict(model_args, **kwargs))
+    model = _create_eva('vit_large_patch16_mrope_ape_224', pretrained=pretrained, **dict(model_args, **kwargs))
     return model