File size: 5,535 Bytes
0522016 | 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 | # Derived from NVIDIA LocateAnything-3B (modeling_vit.py / modeling_locateanything.py, Copyright (c) 2026 NVIDIA
# CORPORATION) under the NVIDIA License (see ../LICENSE); MoonViT is MIT-licensed by Moonshot AI. Modified.
import sys, math, torch, torch.nn as nn, torch.nn.functional as F
from safetensors import safe_open
GH = int(sys.argv[1]) if len(sys.argv) > 1 else 32 # patch grid (H/14, W/14)
GW = int(sys.argv[2]) if len(sys.argv) > 2 else 32
D, HEADS, HD, LAYERS, FF, P = 1152, 16, 72, 27, 4304, 14
MP = "./model"
def load():
w = {}
import json
idx = json.load(open(f"{MP}/model.safetensors.index.json"))["weight_map"]
files = sorted(set(idx.values()))
for f in files:
with safe_open(f"{MP}/{f}", "pt") as sf:
for k in sf.keys():
if k.startswith("vision_model.") or k.startswith("mlp1."):
w[k] = sf.get_tensor(k).float()
return w
def gelu_tanh(x):
return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * x * x * x)))
class Block(nn.Module):
def __init__(s):
super().__init__()
s.norm0 = nn.LayerNorm(D); s.norm1 = nn.LayerNorm(D)
s.wqkv = nn.Linear(D, 3 * D); s.wo = nn.Linear(D, D)
s.fc0 = nn.Linear(D, FF); s.fc1 = nn.Linear(FF, D)
def forward(s, x, cos, sin): # x (1,N,D); cos/sin (1,N,1,36)
N = x.shape[1]
qkv = s.wqkv(s.norm0(x)).view(1, N, 3, HEADS, HD)
q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] # (1,N,H,HD)
def rot(t):
t = t.view(1, N, HEADS, HD // 2, 2)
r, i = t[..., 0], t[..., 1]
return torch.stack([r * cos - i * sin, r * sin + i * cos], dim=-1).view(1, N, HEADS, HD)
q, k = rot(q), rot(k)
q = q.permute(0, 2, 1, 3); k = k.permute(0, 2, 1, 3); v = v.permute(0, 2, 1, 3)
a = torch.softmax(torch.matmul(q, k.transpose(-1, -2)) * (1.0 / math.sqrt(HD)), dim=-1)
o = torch.matmul(a, v).permute(0, 2, 1, 3).reshape(1, N, D)
x = x + s.wo(o)
return x + s.fc1(gelu_tanh(s.fc0(s.norm1(x))))
class ViTFixed(nn.Module):
def __init__(s, w):
super().__init__()
N = GH * GW
s.proj = nn.Conv2d(3, D, P, P)
s.proj.weight.data = w["vision_model.patch_embed.proj.weight"]; s.proj.bias.data = w["vision_model.patch_embed.proj.bias"]
pe = w["vision_model.patch_embed.pos_emb.weight"] # (64,64,D)
if (GH, GW) != tuple(pe.shape[:2]):
pe = F.interpolate(pe.permute(2, 0, 1).unsqueeze(0), size=(GH, GW), mode="bicubic").squeeze(0).permute(1, 2, 0)
s.register_buffer("pos", pe.reshape(1, N, D).contiguous())
# 2D rope: angle[:,2i]=x*f_i, angle[:,2i+1]=y*f_i (x = column, y = row)
f = 1.0 / (10000 ** (torch.arange(0, HD, 4)[: HD // 4].float() / HD))
ys, xs = torch.meshgrid(torch.arange(GH).float(), torch.arange(GW).float(), indexing="ij")
xs, ys = xs.reshape(-1), ys.reshape(-1)
ang = torch.stack([torch.outer(xs, f), torch.outer(ys, f)], dim=-1).reshape(N, HD // 2)
s.register_buffer("cos", ang.cos().view(1, N, 1, HD // 2)); s.register_buffer("sin", ang.sin().view(1, N, 1, HD // 2))
s.blocks = nn.ModuleList([Block() for _ in range(LAYERS)])
for i, b in enumerate(s.blocks):
p = f"vision_model.encoder.blocks.{i}."
b.norm0.weight.data = w[p+"norm0.weight"]; b.norm0.bias.data = w[p+"norm0.bias"]
b.norm1.weight.data = w[p+"norm1.weight"]; b.norm1.bias.data = w[p+"norm1.bias"]
b.wqkv.weight.data = w[p+"wqkv.weight"]; b.wqkv.bias.data = w[p+"wqkv.bias"]
b.wo.weight.data = w[p+"wo.weight"]; b.wo.bias.data = w[p+"wo.bias"]
b.fc0.weight.data = w[p+"mlp.fc0.weight"]; b.fc0.bias.data = w[p+"mlp.fc0.bias"]
b.fc1.weight.data = w[p+"mlp.fc1.weight"]; b.fc1.bias.data = w[p+"mlp.fc1.bias"]
s.fln = nn.LayerNorm(D)
s.fln.weight.data = w["vision_model.encoder.final_layernorm.weight"]; s.fln.bias.data = w["vision_model.encoder.final_layernorm.bias"]
s.m0 = nn.LayerNorm(4 * D); s.m1 = nn.Linear(4 * D, 2048); s.m3 = nn.Linear(2048, 2048)
s.m0.weight.data = w["mlp1.0.weight"]; s.m0.bias.data = w["mlp1.0.bias"]
s.m1.weight.data = w["mlp1.1.weight"]; s.m1.bias.data = w["mlp1.1.bias"]
s.m3.weight.data = w["mlp1.3.weight"]; s.m3.bias.data = w["mlp1.3.bias"]
def forward(s, pixel): # (1,3,H,W), normalized with mean=std=0.5
x = s.proj(pixel).flatten(2).transpose(1, 2) + s.pos # (1,N,D), raster order
for b in s.blocks: x = b(x, s.cos, s.sin)
x = s.fln(x)
x = x.view(GH // 2, 2, GW // 2, 2, D).permute(0, 2, 1, 3, 4).reshape(1, (GH // 2) * (GW // 2), 4 * D)
x = s.m3(F.gelu(s.m1(s.m0(x))))
return x[0] # (N/4, 2048)
if __name__ == "__main__":
m = ViTFixed(load()).eval()
ref = torch.load("ref_vit.pt")
with torch.no_grad(): out = m(ref["pixel"])
r = ref["out"]
cos = F.cosine_similarity(out.flatten(), r.flatten(), dim=0).item()
print("shape", tuple(out.shape), "cos", round(cos, 6), "max|diff|", (out - r).abs().max().item(), "ref std", r.std().item())
if len(sys.argv) > 3 and sys.argv[3] == "export":
torch.onnx.export(m, ref["pixel"], f"out/vit_{GH*14}x{GW*14}.onnx", input_names=["pixel"], output_names=["embeds"], opset_version=17, do_constant_folding=True)
print("exported")
|