Update modeling_vit.py

modeling_vit.py

@@ -25,51 +25,60 @@ class SwiGLU(nn.Module):
 
     def forward(self, x):
         return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))
-
+        
 class RotaryEmbedding2D(nn.Module):
     def __init__(self, head_dim: int, grid_size: int, base: float = 10000.0):
         super().__init__()
         self.head_dim = head_dim
-        self.axis_dim = head_dim // 2
         self.grid_size = grid_size
-        self.num_patches = grid_size * grid_size
+        self.base = base
+        self.axis_dim = head_dim // 2
+        
+        self.cos_sin_cache = None
 
+    def get_cos_sin(self, device, dtype):
+        if self.cos_sin_cache is not None and self.cos_sin_cache[0].device == device:
+            return self.cos_sin_cache
+            
         inv_freq = 1.0 / (
-            base ** (torch.arange(0, self.axis_dim, 2, dtype=torch.float32) / self.axis_dim)
+            self.base ** (torch.arange(0, self.axis_dim, 2, dtype=torch.float32, device=device) / self.axis_dim)
         )
-        coords = torch.arange(grid_size, dtype=torch.float32)
+        
+        coords = torch.arange(self.grid_size, dtype=torch.float32, device=device)
         yy, xx = torch.meshgrid(coords, coords, indexing="ij")
         x_freqs = torch.outer(xx.reshape(-1), inv_freq)
         y_freqs = torch.outer(yy.reshape(-1), inv_freq)
+        
+        cos_x = x_freqs.cos()[None, None, :, :].to(dtype)
+        sin_x = x_freqs.sin()[None, None, :, :].to(dtype)
+        cos_y = y_freqs.cos()[None, None, :, :].to(dtype)
+        sin_y = y_freqs.sin()[None, None, :, :].to(dtype)
+        
+        cos = torch.cat((cos_x, cos_y), dim=-1)
+        sin = torch.cat((sin_x, sin_y), dim=-1)
+        
+        self.cos_sin_cache = (cos, sin)
+        return cos, sin
 
-        self.register_buffer("cos_x", x_freqs.cos()[None, None, :, :], persistent=False)
-        self.register_buffer("sin_x", x_freqs.sin()[None, None, :, :], persistent=False)
-        self.register_buffer("cos_y", y_freqs.cos()[None, None, :, :], persistent=False)
-        self.register_buffer("sin_y", y_freqs.sin()[None, None, :, :], persistent=False)
-
-    @staticmethod
-    def _rotate_axis(x, cos, sin):
+    def apply_rotary_emb(self, x, cos, sin):
         x_even = x[..., 0::2]
         x_odd = x[..., 1::2]
         out_even = x_even * cos - x_odd * sin
         out_odd = x_even * sin + x_odd * cos
         return torch.stack((out_even, out_odd), dim=-1).flatten(-2)
 
-    def _apply_rope(self, x):
-        cls_token = x[:, :, :1, :]
-        patches = x[:, :, 1:, :]
-        x_axis, y_axis = patches.split(self.axis_dim, dim=-1)
-        cos_x = self.cos_x.to(device=x.device, dtype=x.dtype)
-        sin_x = self.sin_x.to(device=x.device, dtype=x.dtype)
-        cos_y = self.cos_y.to(device=x.device, dtype=x.dtype)
-        sin_y = self.sin_y.to(device=x.device, dtype=x.dtype)
-        x_axis = self._rotate_axis(x_axis, cos_x, sin_x)
-        y_axis = self._rotate_axis(y_axis, cos_y, sin_y)
-        patches = torch.cat((x_axis, y_axis), dim=-1)
-        return torch.cat((cls_token, patches), dim=2)
-        
     def forward(self, q, k):
-        return self._apply_rope(q), self._apply_rope(k)
+        cos, sin = self.get_cos_sin(q.device, q.dtype)
+        
+        if q.shape[-2] == cos.shape[-2] + 1:
+            cls_cos = torch.ones(1, 1, 1, cos.shape[-1], device=q.device, dtype=q.dtype)
+            cls_sin = torch.zeros(1, 1, 1, sin.shape[-1], device=q.device, dtype=q.dtype)
+            cos = torch.cat((cls_cos, cos), dim=-2)
+            sin = torch.cat((cls_sin, sin), dim=-2)
+            
+        q_pos = self.apply_rotary_emb(q, cos, sin)
+        k_pos = self.apply_rotary_emb(k, cos, sin)
+        return q_pos, k_pos
 
 class ConvStem(nn.Module):
     def __init__(self, in_chans: int, embed_dim: int, channels: tuple[int, int, int]):