Update modeling_vit.py

modeling_vit.py

@@ -1 +1,223 @@
- 
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import SequenceClassifierOutput
+from .configuration_vit import CustomViTNanoV2Config
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        variance = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(variance + self.eps)
+        return self.weight * x
+
+class SwiGLU(nn.Module):
+    def __init__(self, in_features, hidden_features, out_features):
+        super().__init__()
+        self.w_gate = nn.Linear(in_features, hidden_features, bias=False)
+        self.w_up = nn.Linear(in_features, hidden_features, bias=False)
+        self.w_down = nn.Linear(hidden_features, out_features, bias=False)
+
+    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
+
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, self.axis_dim, 2, dtype=torch.float32) / self.axis_dim)
+        )
+        coords = torch.arange(grid_size, dtype=torch.float32)
+        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)
+
+        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):
+        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)
+
+class ConvStem(nn.Module):
+    def __init__(self, in_chans: int, embed_dim: int, channels: tuple[int, int, int]):
+        super().__init__()
+        c1, c2, c3 = channels
+        self.proj = nn.Sequential(
+            nn.Conv2d(in_chans, c1, kernel_size=3, stride=2, padding=1, bias=False),
+            nn.BatchNorm2d(c1),
+            nn.GELU(),
+            nn.Conv2d(c1, c2, kernel_size=3, stride=2, padding=1, bias=False),
+            nn.BatchNorm2d(c2),
+            nn.GELU(),
+            nn.Conv2d(c2, c3, kernel_size=3, stride=2, padding=1, bias=False),
+            nn.BatchNorm2d(c3),
+            nn.GELU(),
+            nn.Conv2d(c3, embed_dim, kernel_size=3, stride=2, padding=1, bias=False),
+        )
+    def forward(self, x):
+        return self.proj(x)
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads, grid_size, dropout=0.0):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.dropout = float(dropout)
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=False)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(dropout)
+        self.rope = RotaryEmbedding2D(self.head_dim, grid_size=grid_size)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        q, k = self.rope(q, k)
+
+        x = F.scaled_dot_product_attention(
+            q, k, v,
+            dropout_p=(self.dropout if self.training else 0.0),
+            is_causal=False,
+        )
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        return self.proj_drop(x)
+
+class Block(nn.Module):
+    def __init__(self, dim, num_heads, grid_size, mlp_hidden_dim, dropout=0.0):
+        super().__init__()
+        self.norm1 = RMSNorm(dim)
+        self.attn = Attention(dim, num_heads=num_heads, grid_size=grid_size, dropout=dropout)
+        self.norm2 = RMSNorm(dim)
+        self.mlp = nn.Sequential(
+            SwiGLU(dim, mlp_hidden_dim, dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        x = x + self.attn(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+class CustomViTNanoV2PreTrainedModel(PreTrainedModel):
+    config_class = CustomViTNanoV2Config
+    base_model_prefix = "custom_vit"
+    main_input_name = "pixel_values"
+    _no_split_modules = ["Block"]
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Conv2d):
+            nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.BatchNorm2d):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+        elif isinstance(module, RMSNorm):
+            nn.init.ones_(module.weight)
+
+class CustomViTNanoV2ForImageClassification(CustomViTNanoV2PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_classes
+        self.config = config
+        
+        self.patch_size = config.patch_size
+        self.grid_size = config.image_size // config.patch_size
+        
+        self.patch_embed = ConvStem(
+            in_chans=config.in_chans,
+            embed_dim=config.embed_dim,
+            channels=tuple(config.stem_channels),
+        )
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dim))
+        self.pos_drop = nn.Dropout(p=config.dropout)
+        
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    dim=config.embed_dim,
+                    num_heads=config.num_heads,
+                    grid_size=self.grid_size,
+                    mlp_hidden_dim=config.mlp_hidden_dim,
+                    dropout=config.dropout,
+                )
+                for _ in range(config.depth)
+            ]
+        )
+        self.norm = RMSNorm(config.embed_dim)
+        self.head = nn.Linear(config.embed_dim, config.num_classes) if config.num_classes > 0 else nn.Identity()
+        
+        self.post_init()
+
+    def forward(self, pixel_values=None, labels=None, return_dict=None):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        B = pixel_values.shape[0]
+        x = self.patch_embed(pixel_values)
+        x = x.flatten(2).transpose(1, 2)
+        
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = self.pos_drop(x)
+        
+        for block in self.blocks:
+            x = block(x)
+            
+        x = self.norm(x)
+        cls_out = x[:, 0]
+        logits = self.head(cls_out)
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            
+        if not return_dict:
+            output = (logits,)
+            return ((loss,) + output) if loss is not None else output
+            
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+        )
+
+CustomViTNanoV2ForImageClassification.register_for_auto_class("AutoModelForImageClassification")