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")