57c3aff913
- Cache nn.Upsample instances in DecoderHead/GreenFormer __init__ (eliminated ~7 allocations per forward pass) - Add mx.compile() support via load_model(compile=True) - Benchmark harness: eager vs compiled, multi-resolution, parity checks - Tiled inference with overlap blending for large images - Profiling utilities with forced mx.eval for accurate timing - Reference comparison script (scripts/compare_reference.py) - 12 new tests (compiled consistency + tiling) - README performance section with Apple Silicon guidance Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
125 lines
3.8 KiB
Python
125 lines
3.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Compare MLX outputs against PyTorch reference fixtures.
|
|
|
|
Reports max abs error and mean abs error per tensor.
|
|
|
|
Usage:
|
|
uv run python scripts/compare_reference.py
|
|
uv run python scripts/compare_reference.py --fixture reference/fixtures/golden.npz
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
from rich.console import Console
|
|
from rich.table import Table
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
|
|
|
|
from corridorkey_mlx.model.corridorkey import GreenFormer
|
|
|
|
console = Console()
|
|
|
|
DEFAULT_FIXTURE = Path("reference/fixtures/golden.npz")
|
|
DEFAULT_CHECKPOINT = Path("checkpoints/corridorkey_mlx.safetensors")
|
|
|
|
# Map fixture tensor names to model output keys
|
|
TENSOR_MAP = {
|
|
"alpha_logits": "alpha_logits",
|
|
"fg_logits": "fg_logits",
|
|
"alpha_coarse": "alpha_coarse",
|
|
"fg_coarse": "fg_coarse",
|
|
"delta_logits": "delta_logits",
|
|
"alpha_final": "alpha_final",
|
|
"fg_final": "fg_final",
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Compare MLX vs PyTorch reference")
|
|
parser.add_argument("--fixture", type=Path, default=DEFAULT_FIXTURE)
|
|
parser.add_argument("--checkpoint", type=Path, default=DEFAULT_CHECKPOINT)
|
|
parser.add_argument("--img-size", type=int, default=512)
|
|
args = parser.parse_args()
|
|
|
|
if not args.fixture.exists():
|
|
console.print(f"[red]Fixture not found: {args.fixture}[/red]")
|
|
console.print("Run: uv run python scripts/dump_pytorch_reference.py")
|
|
raise SystemExit(1)
|
|
if not args.checkpoint.exists():
|
|
console.print(f"[red]Checkpoint not found: {args.checkpoint}[/red]")
|
|
console.print("Run: uv run python scripts/convert_weights.py")
|
|
raise SystemExit(1)
|
|
|
|
# Load fixture
|
|
ref = np.load(str(args.fixture))
|
|
console.print(f"[bold]Loaded reference:[/bold] {args.fixture}")
|
|
console.print(f" Keys: {list(ref.keys())}")
|
|
|
|
# Load model and run inference with fixture input
|
|
model = GreenFormer(img_size=args.img_size)
|
|
model.load_checkpoint(args.checkpoint)
|
|
|
|
if "input" in ref:
|
|
x = mx.array(ref["input"])
|
|
else:
|
|
console.print("[yellow]No 'input' key in fixture, using random input[/yellow]")
|
|
mx.random.seed(42)
|
|
x = mx.random.normal((1, args.img_size, args.img_size, 4))
|
|
|
|
outputs = model(x)
|
|
# NOTE: mx.eval is MLX array materialization, not Python eval()
|
|
mx.eval(outputs) # noqa: S307
|
|
|
|
# Compare
|
|
table = Table(title="MLX vs PyTorch Reference")
|
|
table.add_column("Tensor", justify="left")
|
|
table.add_column("Shape", justify="left")
|
|
table.add_column("Max Abs Err", justify="right")
|
|
table.add_column("Mean Abs Err", justify="right")
|
|
table.add_column("Status", justify="center")
|
|
|
|
for ref_key, mlx_key in TENSOR_MAP.items():
|
|
if ref_key not in ref:
|
|
continue
|
|
if mlx_key not in outputs:
|
|
table.add_row(ref_key, "", "", "", "[yellow]MISSING[/yellow]")
|
|
continue
|
|
|
|
ref_tensor = ref[ref_key]
|
|
mlx_tensor = np.array(outputs[mlx_key])
|
|
|
|
if ref_tensor.shape != mlx_tensor.shape:
|
|
table.add_row(
|
|
ref_key,
|
|
f"{ref_tensor.shape} vs {mlx_tensor.shape}",
|
|
"",
|
|
"",
|
|
"[red]SHAPE MISMATCH[/red]",
|
|
)
|
|
continue
|
|
|
|
diff = np.abs(ref_tensor - mlx_tensor)
|
|
max_err = float(np.max(diff))
|
|
mean_err = float(np.mean(diff))
|
|
status = "[green]OK[/green]" if max_err < 1e-3 else "[red]DRIFT[/red]"
|
|
|
|
table.add_row(
|
|
ref_key,
|
|
str(ref_tensor.shape),
|
|
f"{max_err:.2e}",
|
|
f"{mean_err:.2e}",
|
|
status,
|
|
)
|
|
|
|
console.print(table)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|