Views
No views yet
MoonVit.py script (build / eval / inference).eval below).| File | Description |
|---|---|
moonvit.onnx | The exported graph. Legacy export = single file (~1.7 GB); dynamo export = graph + moonvit.onnx.data (external weights). |
manifest.json | Export metadata — grid, I/O shapes, normalization, which exporter was used. |
README.md | This file. |
grid_hws.tolist() (per-image
Python loops for the interpolated position embedding, the 2D RoPE, and the 2×2 patch merger), so
the graph is data-dependent on the image grid. This export bakes one fixed grid as a
constant. The default is 28×28 patches = 392×392 px; check manifest.json for the exact grid
this file was built with. For other input resolutions, export another variant
(--grid-h/--grid-w).pixel_values: float32 [L, 3, 14, 14] — L = grid_h × grid_w packed 14×14 patches,
raster order, normalized with mean/std = 0.5.image_features: float32 [L/4, 4, 1152] — merged tokens after the 2×2 patch merger
(e.g. 28×28 → [196, 4, 1152]).1import json, numpy as np, onnxruntime as ort
2from PIL import Image
3
4d = "." # this folder
5man = json.load(open(f"{d}/manifest.json"))
6gh, gw, P = man["grid_h"], man["grid_w"], man["patch_size"]
7
8# preprocess an image to the export grid (raster-order packed patches, mean/std 0.5)
9img = Image.open("figures/demo.png").convert("RGB").resize((gw * P, gh * P), Image.BICUBIC)
10x = (np.asarray(img, np.float32) / 255.0 - 0.5) / 0.5 # [H,W,3]
11x = x.transpose(2, 0, 1).reshape(3, gh, P, gw, P).transpose(1, 3, 0, 2, 4)
12pixel_values = np.ascontiguousarray(x.reshape(gh * gw, 3, P, P))
13
14sess = ort.InferenceSession(f"{d}/{man['file']}", providers=["CPUExecutionProvider"])
15feats = sess.run(None, {"pixel_values": pixel_values})[0] # [gh*gw/4, 4, 1152]
16print(feats.shape)uv run MoonVit.py inference --onnx-model-path <this_dir> --image img.png.MoonVit.py)1uv run MoonVit.py build --output MoonVitOnnx # legacy exporter (default)
2uv run MoonVit.py build --output MoonVitOnnx --dynamo # new torch.onnx dynamo exporter
3uv run MoonVit.py eval --onnx-model-path MoonVitOnnx # original PyTorch vs ONNX
4uv run MoonVit.py inference --onnx-model-path MoonVitOnnx --image img.png| Exporter | Graph | Weights | Notes |
|---|---|---|---|
legacy (default) | ~5,050 nodes | inline (~1.7 GB single file) | unrolls the data-dependent loops against the baked grid |
dynamo (--dynamo) | ~1,794 nodes | external .onnx.data | cleaner/smaller graph; requires onnxscript |
freqs_cis is
precomputed for the baked grid and applied as real-valued rotation.PytorchGELUTanh shim — the upstream remote code imports it from transformers.activations
(removed in newer transformers); it's exactly nn.GELU(approximate="tanh").Slice/Gather index tensors.grid=28x28 original PyTorch vs ONNX
sample 0: cos=1.000000 max|Δ|=9.995e-02
sample 1: cos=1.000000 max|Δ|=6.195e-03
sample 2: cos=1.000000 max|Δ|=2.052e-03
=== PASS (worst cosine 1.000000, tol 0.999) ===max|Δ| is on large-magnitude features, feature std ≈ 4.2 — cosine 1.0 is the real signal.)1from PIL import Image
2from transformers import AutoModel, AutoImageProcessor
3
4model_path = "moonshotai/MoonViT-SO-400M"
5model = AutoModel.from_pretrained(model_path, torch_dtype="auto", device_map="auto",
6 trust_remote_code=True)
7processor = AutoImageProcessor.from_pretrained(model_path, trust_remote_code=True)
8
9image = Image.open("./figures/demo.png")
10proc = processor(image, return_tensors="pt").to(dtype=model.dtype, device=model.device)
11image_features: list = model(proc.pixel_values, proc.image_grid_hws)
12print(image_features[0].dtype, image_features[0].shape) # e.g. bf16, [N, 4, 1152]