File size: 4,460 Bytes
fa50b55 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | from functools import lru_cache
import torch
import torch.nn.functional as F
from torch import nn
@lru_cache(8)
def get_vision_cos_sin(n_h: int, n_w: int, dim: int, theta: float):
inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
hpos = torch.arange(n_h).unsqueeze(1).expand(n_h, n_w)
wpos = torch.arange(n_w).unsqueeze(0).expand(n_h, n_w)
freqs = torch.stack([hpos, wpos], dim=-1).reshape(-1, 2, 1).float() * inv_freq
freqs = freqs.flatten(1)
return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1)
def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
dtype = x.dtype
x1, x2 = x.float().chunk(2, dim=-1)
return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1).to(dtype)
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, dtype=torch.float32))
def forward(self, x: torch.Tensor) -> torch.Tensor:
dtype = x.dtype
x = x.float()
x = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps)
return (self.weight * x).to(dtype)
class PatchEmbed(nn.Module):
def __init__(self, args):
super().__init__()
self.proj = nn.Linear(3 * args.vision_patch_size ** 2, args.vision_dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.proj(x.flatten(1))
class Attention(nn.Module):
def __init__(self, args):
super().__init__()
self.n_heads = args.vision_n_heads
self.head_dim = args.vision_dim // args.vision_n_heads
self.wqkv = nn.Linear(args.vision_dim, 3 * args.vision_dim)
self.wo = nn.Linear(args.vision_dim, args.vision_dim)
def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
n = x.size(0)
q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in self.wqkv(x).chunk(3, dim=-1))
q = apply_rotary(q, cos, sin)
k = apply_rotary(k, cos, sin)
o = F.scaled_dot_product_attention(q.transpose(0, 1), k.transpose(0, 1), v.transpose(0, 1))
return self.wo(o.transpose(0, 1).reshape(n, -1))
class MLP(nn.Module):
def __init__(self, args):
super().__init__()
self.w1 = nn.Linear(args.vision_dim, 2 * args.vision_inter_dim, bias=False)
self.w2 = nn.Linear(args.vision_inter_dim, args.vision_dim, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
gate, up = self.w1(x).chunk(2, dim=-1)
return self.w2(F.silu(gate) * up)
class Block(nn.Module):
def __init__(self, args):
super().__init__()
self.norm1 = RMSNorm(args.vision_dim)
self.attn = Attention(args)
self.norm2 = RMSNorm(args.vision_dim)
self.mlp = MLP(args)
def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
x = x + self.attn(self.norm1(x), cos, sin)
return x + self.mlp(self.norm2(x))
class ViT(nn.Module):
"""DeepSeek ViT: full bidirectional attention over one image with 2D RoPE."""
def __init__(self, args):
super().__init__()
self.rope_dim = args.vision_dim // args.vision_n_heads // 2
self.rope_theta = args.vision_rope_theta
self.patch_embed = PatchEmbed(args)
self.blocks = nn.ModuleList([Block(args) for _ in range(args.vision_n_layers)])
self.norm = RMSNorm(args.vision_dim)
def forward(self, patches: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor:
x = self.patch_embed(patches)
cos, sin = get_vision_cos_sin(n_h, n_w, self.rope_dim, self.rope_theta)
for block in self.blocks:
x = block(x, cos, sin)
return self.norm(x)
class Aligner(nn.Module):
def __init__(self, args):
super().__init__()
self.downsample_ratio = args.vision_downsample_ratio
in_dim = args.vision_dim * self.downsample_ratio ** 2
self.w1 = nn.Linear(in_dim, args.dim)
self.w2 = nn.Linear(args.dim, args.dim)
def forward(self, x: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor:
r = self.downsample_ratio
x = x.view(n_h, n_w, -1).permute(2, 0, 1)
x = F.pad(x, (0, -n_w % r, 0, -n_h % r))
x = F.unfold(x.unsqueeze(0), r, stride=r).squeeze(0).transpose(0, 1)
return self.w2(F.gelu(self.w1(x)))
|