modeling_vit.py

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 CustomViTNanoConfig

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.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 / (
            self.base ** (torch.arange(0, self.axis_dim, 2, dtype=torch.float32, device=device) / self.axis_dim)
        )
        
        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

    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 forward(self, q, 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]):
        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 CustomViTNanoPreTrainedModel(PreTrainedModel):
    config_class = CustomViTNanoConfig
    base_model_prefix = ""
    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 CustomViTNanoForImageClassification(CustomViTNanoPreTrainedModel):
    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,
        )

CustomViTNanoForImageClassification.register_for_auto_class("AutoModelForImageClassification")