|
|
| """Dump a DINOv3 reference forward pass for `cargo run --example verify`.
|
|
|
| Writes three files into the output directory:
|
|
|
| pixel_values.bin f32 [3, S, S] the preprocessed input tensor
|
| features.bin f32 [tokens, 384] expected last_hidden_state
|
| model.safetensors weights, in the graph's naming
|
|
|
| Dumping the *preprocessed* pixel tensor rather than an image is deliberate:
|
| it keeps resize and normalization differences out of the comparison, so a
|
| mismatch in `verify` is a mismatch in the graph, not in the resampling
|
| filter.
|
|
|
| Saving the state dict here rather than downloading it in Rust also means
|
| `verify` needs no Hub access and no license acceptance at run time.
|
|
|
| pip install torch transformers safetensors pillow numpy
|
| huggingface-cli login # facebook/... is license-gated
|
| python tools/dump_reference.py --out ref/
|
|
|
| The weights are gated. Accept the DINOv3 license on the canonical model page |
| and authenticate with Hugging Face before running this script. |
| """
|
|
|
| import argparse |
| import hashlib |
| import json |
| import pathlib |
|
|
| import numpy as np |
| import torch |
| import transformers |
| from safetensors.torch import save_file |
| from transformers import AutoModel |
| from transformers.models.dinov3_vit.modeling_dinov3_vit import ( |
| apply_rotary_pos_emb, |
| eager_attention_forward, |
| ) |
|
|
| DEFAULT_MODEL = "facebook/dinov3-vits16-pretrain-lvd1689m"
|
| IMAGE_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
| IMAGE_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
|
|
|
|
| def build_input(size: int, image: str | None) -> np.ndarray:
|
| """Return a normalized f32 CHW tensor."""
|
| if image is not None:
|
| from PIL import Image
|
|
|
| rgb = Image.open(image).convert("RGB").resize((size, size), Image.BILINEAR)
|
| hwc = np.asarray(rgb, dtype=np.float32) / 255.0
|
| else:
|
|
|
|
|
|
|
| ys, xs = np.mgrid[0:size, 0:size].astype(np.float32)
|
| u, v = xs / size, ys / size
|
| hwc = np.stack(
|
| [
|
| 0.5 + 0.5 * np.sin(6.0 * u + 2.0 * v),
|
| 0.5 + 0.5 * np.sin(4.0 * v - 3.0 * u * v),
|
| 0.5 + 0.5 * np.cos(5.0 * u * v + u),
|
| ],
|
| axis=-1,
|
| ).astype(np.float32)
|
|
|
| chw = ((hwc - IMAGE_MEAN) / IMAGE_STD).transpose(2, 0, 1)
|
| return np.ascontiguousarray(chw, dtype=np.float32)
|
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser()
|
| ap.add_argument("--out", default="ref", type=pathlib.Path)
|
| ap.add_argument("--model", default=DEFAULT_MODEL) |
| ap.add_argument("--size", type=int, default=224) |
| ap.add_argument( |
| "--layers", |
| type=int, |
| default=12, |
| help="number of leading encoder layers to execute before the final norm", |
| ) |
| ap.add_argument("--image", default=None, help="optional real photo instead of a test pattern")
|
| args = ap.parse_args()
|
|
|
| if args.size % 16:
|
| raise SystemExit(f"--size {args.size} is not a multiple of the patch size (16)")
|
|
|
| args.out.mkdir(parents=True, exist_ok=True)
|
|
|
| model = AutoModel.from_pretrained(args.model, dtype=torch.float32).eval()
|
| cfg = model.config
|
| print(
|
| f"{args.model}: {cfg.num_hidden_layers} layers, hidden {cfg.hidden_size}, "
|
| f"{cfg.num_register_tokens} register tokens, gated_mlp={cfg.use_gated_mlp}"
|
| )
|
| if cfg.use_gated_mlp: |
| raise SystemExit("the Rust graph implements the plain MLP only (ViT-S/B)") |
| if not 0 <= args.layers <= cfg.num_hidden_layers: |
| raise SystemExit( |
| f"--layers must be in [0, {cfg.num_hidden_layers}], got {args.layers}" |
| ) |
|
|
| chw = build_input(args.size, args.image) |
| pixels = torch.from_numpy(chw)[None] |
| with torch.no_grad(): |
| hidden_states = model.embeddings(pixels) |
| embeddings = hidden_states[0].numpy().astype(np.float32) |
| position_embeddings = model.rope_embeddings(pixels) |
| first_layer = model.model.layer[0] |
| first_norm1 = first_layer.norm1(hidden_states) |
| first_q = first_layer.attention.q_proj(first_norm1) |
| first_k = first_layer.attention.k_proj(first_norm1) |
| first_v = first_layer.attention.v_proj(first_norm1) |
| tokens = first_q.shape[1] |
| heads = cfg.num_attention_heads |
| head_dim = cfg.hidden_size // heads |
| q_heads = first_q.view(1, tokens, heads, head_dim).transpose(1, 2) |
| k_heads = first_k.view(1, tokens, heads, head_dim).transpose(1, 2) |
| v_heads = first_v.view(1, tokens, heads, head_dim).transpose(1, 2) |
| q_rope, k_rope = apply_rotary_pos_emb( |
| q_heads, k_heads, *position_embeddings |
| ) |
| first_attention, _ = eager_attention_forward( |
| first_layer.attention, |
| q_rope, |
| k_rope, |
| v_heads, |
| None, |
| scaling=first_layer.attention.scaling, |
| ) |
| first_attention = first_attention.reshape(1, tokens, cfg.hidden_size) |
| first_attention_projected = first_layer.attention.o_proj(first_attention) |
| first_attention_scaled = first_layer.layer_scale1( |
| first_attention_projected |
| ) |
| first_residual = hidden_states + first_attention_scaled |
| first_norm2 = first_layer.norm2(first_residual) |
| first_mlp_up = first_layer.mlp.up_proj(first_norm2) |
| first_mlp_activated = first_layer.mlp.act_fn(first_mlp_up) |
| first_mlp_down = first_layer.mlp.down_proj(first_mlp_activated) |
| first_mlp_scaled = first_layer.layer_scale2(first_mlp_down) |
| first_output = first_residual + first_mlp_scaled |
| first_final_norm = model.norm(first_output) |
| for layer in model.model.layer[: args.layers]: |
| hidden_states = layer( |
| hidden_states, position_embeddings=position_embeddings |
| ) |
| features = model.norm(hidden_states)[0].numpy().astype(np.float32) |
|
|
| grid = args.size // cfg.patch_size
|
| expected_tokens = 1 + cfg.num_register_tokens + grid * grid
|
| assert features.shape == (expected_tokens, cfg.hidden_size), features.shape
|
|
|
| (args.out / "pixel_values.bin").write_bytes(chw.tobytes()) |
| (args.out / "embeddings.bin").write_bytes( |
| np.ascontiguousarray(embeddings).tobytes() |
| ) |
| for name, tensor in { |
| "first-norm1.bin": first_norm1, |
| "first-q.bin": first_q, |
| "first-k.bin": first_k, |
| "first-v.bin": first_v, |
| "first-q-rope.bin": q_rope.transpose(1, 2).reshape( |
| 1, tokens, cfg.hidden_size |
| ), |
| "first-k-rope.bin": k_rope.transpose(1, 2).reshape( |
| 1, tokens, cfg.hidden_size |
| ), |
| "first-attention.bin": first_attention, |
| "first-attention-projected.bin": first_attention_projected, |
| "first-attention-scaled.bin": first_attention_scaled, |
| "first-residual.bin": first_residual, |
| "first-norm2.bin": first_norm2, |
| "first-mlp-up.bin": first_mlp_up, |
| "first-mlp-activated.bin": first_mlp_activated, |
| "first-mlp-down.bin": first_mlp_down, |
| "first-mlp-scaled.bin": first_mlp_scaled, |
| "first-output.bin": first_output, |
| "first-final-norm.bin": first_final_norm, |
| }.items(): |
| (args.out / name).write_bytes( |
| np.ascontiguousarray(tensor[0].numpy().astype(np.float32)).tobytes() |
| ) |
| (args.out / "features.bin").write_bytes(np.ascontiguousarray(features).tobytes()) |
|
|
|
|
| state = {k: v.contiguous().clone() for k, v in model.state_dict().items()}
|
| save_file(state, str(args.out / "model.safetensors")) |
|
|
| def sha256(path: pathlib.Path) -> str: |
| return hashlib.sha256(path.read_bytes()).hexdigest() |
|
|
| local_model = pathlib.Path(args.model) |
| reference = { |
| "schema_version": 1, |
| "base_model": DEFAULT_MODEL if local_model.is_dir() else args.model, |
| "model_source": str(args.model), |
| "image_size": args.size, |
| "encoder_layers": args.layers, |
| "tokens": expected_tokens, |
| "hidden_size": cfg.hidden_size, |
| "torch_version": torch.__version__, |
| "transformers_version": transformers.__version__, |
| "numpy_version": np.__version__, |
| "pixel_values_sha256": sha256(args.out / "pixel_values.bin"), |
| "embeddings_sha256": sha256(args.out / "embeddings.bin"), |
| "features_sha256": sha256(args.out / "features.bin"), |
| "exported_model_sha256": sha256(args.out / "model.safetensors"), |
| } |
| if local_model.is_dir() and (local_model / "model.safetensors").is_file(): |
| reference["source_model_sha256"] = sha256( |
| local_model / "model.safetensors" |
| ) |
| (args.out / "reference.json").write_text( |
| json.dumps(reference, indent=2) + "\n", encoding="utf-8" |
| ) |
|
|
| print( |
| f"wrote {args.out}/ — {args.layers} layers, {args.size}x{args.size}, " |
| f"{grid}x{grid} grid, {expected_tokens} tokens" |
| ) |
| print(f"features: mean {features.mean():+.4f} std {features.std():.4f}")
|
| print(f"\nnow run: cargo run --release --example verify -- {args.out}")
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
