refactor(tests): restructure suite — 15 files to 8, consolidate parity
- Add conftest.py: shared paths, tolerances, skip markers, fixtures - Consolidate imports/shapes/smoke/forward → test_model_contract.py - Consolidate all parity (decoder/refiner/backbone/e2e) → test_parity.py - Fold 2048 slow test into test_engine.py - Rework test_conversion.py to use public convert_checkpoint() API - Simplify test_weights.py: drop CLI/env var tests, keep checksum+config - Rename tiling/compilation test files for consistency - Delete 10 redundant files 80 tests collected (75 pass, 4 skip, 1 deselected) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
cmoyates
parent
62645a2e93
commit
838bd0dd39
71
tests/conftest.py
Normal file
71
tests/conftest.py
Normal file
@ -0,0 +1,71 @@
|
||||
"""Shared test fixtures, paths, tolerances, and skip markers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from corridorkey_mlx.model.corridorkey import GreenFormer
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Paths
|
||||
# ---------------------------------------------------------------------------
|
||||
GOLDEN_PATH = Path("reference/fixtures/golden.npz")
|
||||
GOLDEN_WEIGHTS_PATH = Path("reference/fixtures/golden_weights.npz")
|
||||
MLX_CHECKPOINT_PATH = Path("checkpoints/corridorkey_mlx.safetensors")
|
||||
PT_CHECKPOINT_PATH = Path("checkpoints/CorridorKey_v1.0.pth")
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tolerances
|
||||
# ---------------------------------------------------------------------------
|
||||
PARITY_TOL_TIGHT = 1e-4
|
||||
PARITY_TOL_BACKBONE = 2e-2
|
||||
PARITY_TOL_E2E = 1e-3
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Skip markers
|
||||
# ---------------------------------------------------------------------------
|
||||
has_golden = pytest.mark.skipif(not GOLDEN_PATH.exists(), reason="golden.npz not found")
|
||||
has_checkpoint = pytest.mark.skipif(
|
||||
not MLX_CHECKPOINT_PATH.exists(), reason="MLX checkpoint not found"
|
||||
)
|
||||
has_pt_checkpoint = pytest.mark.skipif(
|
||||
not PT_CHECKPOINT_PATH.exists(), reason="PyTorch checkpoint not found"
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shared fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
IMG_SIZE = 512
|
||||
SMALL_IMG_SIZE = 256
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def golden_fixtures() -> dict[str, np.ndarray]:
|
||||
"""Load golden.npz fixture file (PyTorch NCHW format)."""
|
||||
if not GOLDEN_PATH.exists():
|
||||
pytest.skip("golden.npz not found")
|
||||
return dict(np.load(GOLDEN_PATH))
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def random_model() -> GreenFormer:
|
||||
"""GreenFormer with random weights at 512."""
|
||||
model = GreenFormer(img_size=IMG_SIZE)
|
||||
model.eval()
|
||||
# mx.eval is MLX array materialization, not Python eval()
|
||||
mx.eval(model.parameters()) # noqa: S307
|
||||
return model
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def loaded_model() -> GreenFormer:
|
||||
"""GreenFormer with checkpoint weights at 512."""
|
||||
if not MLX_CHECKPOINT_PATH.exists():
|
||||
pytest.skip("MLX checkpoint not found")
|
||||
model = GreenFormer(img_size=IMG_SIZE)
|
||||
model.load_checkpoint(MLX_CHECKPOINT_PATH)
|
||||
return model
|
||||
@ -1,185 +1,83 @@
|
||||
"""Tests: weight conversion PyTorch → MLX (Phase 3).
|
||||
"""Weight conversion tests through public API.
|
||||
|
||||
Validates key mapping, shape transforms, and round-trip integrity.
|
||||
Reworked from internal convert_state_dict tests to use convert_checkpoint().
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from corridorkey_mlx.convert.converter import convert_checkpoint
|
||||
from corridorkey_mlx.model.corridorkey import GreenFormer
|
||||
|
||||
from .conftest import PT_CHECKPOINT_PATH, has_pt_checkpoint
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections import OrderedDict
|
||||
from pathlib import Path
|
||||
|
||||
from corridorkey_mlx.convert.converter import (
|
||||
REFINER_STEM_MAP,
|
||||
SKIP_SUFFIXES,
|
||||
convert_state_dict,
|
||||
load_pytorch_checkpoint,
|
||||
)
|
||||
EXPECTED_KEY_COUNT = 365
|
||||
|
||||
CHECKPOINT_PATH = Path("checkpoints/CorridorKey_v1.0.pth")
|
||||
SAFETENSORS_PATH = Path("checkpoints/corridorkey_mlx.safetensors")
|
||||
|
||||
def _torch_available() -> bool:
|
||||
try:
|
||||
import torch # noqa: F401
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
has_torch = pytest.mark.skipif(not _torch_available(), reason="torch not installed")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def checkpoint_data() -> tuple[dict[str, np.ndarray], OrderedDict[str, np.ndarray], list]:
|
||||
"""Load and convert checkpoint once for all tests in this module."""
|
||||
if not CHECKPOINT_PATH.exists():
|
||||
pytest.skip("Checkpoint not found — need CorridorKey_v1.0.pth")
|
||||
state_dict = load_pytorch_checkpoint(CHECKPOINT_PATH)
|
||||
converted, diagnostics = convert_state_dict(state_dict)
|
||||
return state_dict, converted, diagnostics
|
||||
def converted_path(tmp_path_factory: pytest.TempPathFactory) -> Path:
|
||||
"""Convert PT checkpoint to safetensors in temp dir."""
|
||||
if not PT_CHECKPOINT_PATH.exists():
|
||||
pytest.skip("PyTorch checkpoint not found")
|
||||
if not _torch_available():
|
||||
pytest.skip("torch not installed")
|
||||
output = tmp_path_factory.mktemp("weights") / "converted.safetensors"
|
||||
convert_checkpoint(PT_CHECKPOINT_PATH, output)
|
||||
return output
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Key mapping completeness
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestKeyMapping:
|
||||
"""Every PyTorch key maps to an MLX key with no orphans."""
|
||||
|
||||
def test_no_orphan_source_keys(self, checkpoint_data) -> None:
|
||||
"""All source keys are either mapped or explicitly skipped."""
|
||||
state_dict, _, diagnostics = checkpoint_data
|
||||
|
||||
mapped_src_keys = {r.source_key for r in diagnostics}
|
||||
skipped_keys = {k for k in state_dict if any(k.endswith(s) for s in SKIP_SUFFIXES)}
|
||||
accounted_for = mapped_src_keys | skipped_keys
|
||||
|
||||
orphans = set(state_dict.keys()) - accounted_for
|
||||
assert orphans == set(), f"Orphan source keys: {orphans}"
|
||||
|
||||
def test_no_duplicate_dest_keys(self, checkpoint_data) -> None:
|
||||
"""No two source keys map to the same destination key."""
|
||||
_, _, diagnostics = checkpoint_data
|
||||
|
||||
dest_keys = [r.dest_key for r in diagnostics]
|
||||
assert len(dest_keys) == len(set(dest_keys)), "Duplicate destination keys found"
|
||||
|
||||
def test_refiner_stem_remapped(self, checkpoint_data) -> None:
|
||||
"""Refiner stem Sequential keys correctly remapped to named attrs."""
|
||||
_, converted, _ = checkpoint_data
|
||||
|
||||
for pt_key, mlx_key in REFINER_STEM_MAP.items():
|
||||
assert mlx_key in converted, f"Expected remapped key {mlx_key} not found"
|
||||
assert pt_key not in converted, f"Original key {pt_key} should not be in output"
|
||||
|
||||
def test_num_batches_tracked_dropped(self, checkpoint_data) -> None:
|
||||
"""BatchNorm num_batches_tracked keys are not in output."""
|
||||
_, converted, _ = checkpoint_data
|
||||
|
||||
for key in converted:
|
||||
assert "num_batches_tracked" not in key, f"Found num_batches_tracked: {key}"
|
||||
|
||||
def test_expected_key_count(self, checkpoint_data) -> None:
|
||||
"""Output has expected number of keys (367 source - 2 skipped = 365)."""
|
||||
state_dict, converted, _ = checkpoint_data
|
||||
|
||||
source_count = len(state_dict)
|
||||
skip_count = sum(1 for k in state_dict if any(k.endswith(s) for s in SKIP_SUFFIXES))
|
||||
expected = source_count - skip_count
|
||||
assert len(converted) == expected, f"Expected {expected} keys, got {len(converted)}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Conv weight transpose
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestConvWeightTranspose:
|
||||
"""Conv weights correctly transposed from NCHW → NHWC."""
|
||||
|
||||
def test_patch_embed_conv_shape(self, checkpoint_data) -> None:
|
||||
"""Patch embed conv transposed: (112,4,7,7) → (112,7,7,4)."""
|
||||
_, converted, _ = checkpoint_data
|
||||
|
||||
key = "encoder.model.patch_embed.proj.weight"
|
||||
assert key in converted
|
||||
assert converted[key].shape == (112, 7, 7, 4)
|
||||
|
||||
def test_decoder_conv_shapes(self, checkpoint_data) -> None:
|
||||
"""Decoder conv weights transposed correctly."""
|
||||
_, converted, _ = checkpoint_data
|
||||
|
||||
# linear_fuse: (256, 1024, 1, 1) → (256, 1, 1, 1024)
|
||||
for prefix in ("alpha_decoder", "fg_decoder"):
|
||||
fuse_key = f"{prefix}.linear_fuse.weight"
|
||||
assert converted[fuse_key].shape == (256, 1, 1, 1024)
|
||||
|
||||
def test_refiner_conv_shapes(self, checkpoint_data) -> None:
|
||||
"""Refiner conv weights transposed correctly."""
|
||||
_, converted, _ = checkpoint_data
|
||||
|
||||
# stem_conv: (64, 7, 3, 3) → (64, 3, 3, 7)
|
||||
assert converted["refiner.stem_conv.weight"].shape == (64, 3, 3, 7)
|
||||
|
||||
# res block convs: (64, 64, 3, 3) → (64, 3, 3, 64)
|
||||
for i in range(1, 5):
|
||||
for j in range(1, 3):
|
||||
key = f"refiner.res{i}.conv{j}.weight"
|
||||
assert converted[key].shape == (64, 3, 3, 64), f"Wrong shape for {key}"
|
||||
|
||||
# final: (4, 64, 1, 1) → (4, 1, 1, 64)
|
||||
assert converted["refiner.final.weight"].shape == (4, 1, 1, 64)
|
||||
|
||||
def test_linear_weights_unchanged(self, checkpoint_data) -> None:
|
||||
"""Linear (2D) weights are not transposed."""
|
||||
_, _, diagnostics = checkpoint_data
|
||||
|
||||
passthrough_2d = [
|
||||
r for r in diagnostics if r.transform == "passthrough" and len(r.source_shape) == 2
|
||||
]
|
||||
for record in passthrough_2d:
|
||||
assert record.source_shape == record.dest_shape, (
|
||||
f"{record.source_key}: shape changed despite passthrough"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4-channel first conv
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestFirstConv4Channel:
|
||||
"""Patched 4-channel first conv preserved during conversion."""
|
||||
|
||||
def test_input_channels_preserved(self, checkpoint_data) -> None:
|
||||
"""Patch embed conv has 4 input channels (RGB + alpha hint)."""
|
||||
_, converted, _ = checkpoint_data
|
||||
|
||||
weight = converted["encoder.model.patch_embed.proj.weight"]
|
||||
# MLX layout: (O, H, W, I) — I is last dim
|
||||
input_channels = weight.shape[-1]
|
||||
assert input_channels == 4, f"Expected 4 input channels, got {input_channels}"
|
||||
|
||||
def test_values_preserved(self, checkpoint_data) -> None:
|
||||
"""Conv values are preserved (only transposed, not modified)."""
|
||||
state_dict, converted, _ = checkpoint_data
|
||||
|
||||
pt_weight = state_dict["encoder.model.patch_embed.proj.weight"]
|
||||
mlx_weight = converted["encoder.model.patch_embed.proj.weight"]
|
||||
|
||||
# Transpose back to PyTorch layout and compare
|
||||
roundtrip = np.transpose(mlx_weight, (0, 3, 1, 2))
|
||||
np.testing.assert_array_equal(roundtrip, pt_weight)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Safetensors output validation
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestSafetensorsOutput:
|
||||
"""Validate saved safetensors file matches conversion output."""
|
||||
|
||||
def test_safetensors_loadable(self, checkpoint_data) -> None:
|
||||
"""Saved safetensors file can be loaded and has correct keys."""
|
||||
if not SAFETENSORS_PATH.exists():
|
||||
pytest.skip("Run convert_weights.py first")
|
||||
@has_pt_checkpoint
|
||||
@has_torch
|
||||
class TestConvertCheckpoint:
|
||||
def test_output_file_exists(self, converted_path: Path) -> None:
|
||||
assert converted_path.exists()
|
||||
|
||||
def test_safetensors_loadable(self, converted_path: Path) -> None:
|
||||
from safetensors.numpy import load_file
|
||||
|
||||
loaded = load_file(str(SAFETENSORS_PATH))
|
||||
weights = load_file(str(converted_path))
|
||||
assert len(weights) > 0
|
||||
|
||||
_, converted, _ = checkpoint_data
|
||||
def test_key_count(self, converted_path: Path) -> None:
|
||||
from safetensors.numpy import load_file
|
||||
|
||||
assert set(loaded.keys()) == set(converted.keys())
|
||||
for key in converted:
|
||||
np.testing.assert_array_equal(loaded[key], converted[key])
|
||||
weights = load_file(str(converted_path))
|
||||
assert len(weights) == EXPECTED_KEY_COUNT
|
||||
|
||||
def test_model_loads_and_runs(self, converted_path: Path) -> None:
|
||||
"""Roundtrip: convert -> load into GreenFormer -> forward succeeds."""
|
||||
import mlx.core as mx
|
||||
|
||||
model = GreenFormer(img_size=256)
|
||||
model.load_checkpoint(converted_path)
|
||||
|
||||
x = mx.random.normal((1, 256, 256, 4))
|
||||
out = model(x)
|
||||
# mx.eval is MLX array materialization, not Python eval()
|
||||
mx.eval(out) # noqa: S307
|
||||
|
||||
assert out["alpha_final"].shape == (1, 256, 256, 1)
|
||||
assert out["fg_final"].shape == (1, 256, 256, 3)
|
||||
|
||||
# No NaN
|
||||
alpha = np.array(out["alpha_final"])
|
||||
assert not np.isnan(alpha).any()
|
||||
|
||||
@ -1,150 +0,0 @@
|
||||
"""Parity tests: MLX decoder heads vs PyTorch reference (Phase 2).
|
||||
|
||||
Uses saved backbone features -> runs MLX decoder -> compares
|
||||
against saved PyTorch coarse predictions.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from corridorkey_mlx.model.decoder import DecoderHead
|
||||
from corridorkey_mlx.utils.layout import conv_weight_pt_to_mlx, nchw_to_nhwc_np, nhwc_to_nchw_np
|
||||
|
||||
FIXTURE_PATH = Path("reference/fixtures/golden.npz")
|
||||
WEIGHTS_PATH = Path("reference/fixtures/golden_weights.npz")
|
||||
|
||||
BACKBONE_CHANNELS = [112, 224, 448, 896]
|
||||
EMBED_DIM = 256
|
||||
MAX_ABS_TOL = 1e-4 # relaxed slightly for float32 Metal vs CPU differences
|
||||
|
||||
|
||||
def _skip_if_missing() -> None:
|
||||
if not FIXTURE_PATH.exists() or not WEIGHTS_PATH.exists():
|
||||
pytest.skip("Fixture files not found — run dump_pytorch_reference.py first")
|
||||
|
||||
|
||||
def _load_decoder_weights(
|
||||
prefix: str,
|
||||
weights: dict[str, np.ndarray],
|
||||
) -> dict[str, mx.array]:
|
||||
"""Convert PyTorch decoder state_dict to MLX parameter dict.
|
||||
|
||||
Maps PyTorch keys (e.g. 'alpha_decoder.linear_c1.proj.weight') to
|
||||
MLX nested keys (e.g. 'linear_c1.proj.weight'), transposing conv weights.
|
||||
"""
|
||||
params: dict[str, mx.array] = {}
|
||||
for pt_key, value in weights.items():
|
||||
if not pt_key.startswith(prefix):
|
||||
continue
|
||||
mlx_key = pt_key[len(prefix) :]
|
||||
|
||||
# Conv2d weights: (O,I,H,W) -> (O,H,W,I)
|
||||
if "linear_fuse.weight" in mlx_key or "classifier.weight" in mlx_key:
|
||||
value = conv_weight_pt_to_mlx(value)
|
||||
|
||||
# BatchNorm: running_mean -> running_mean, running_var -> running_var
|
||||
# PyTorch uses num_batches_tracked which MLX doesn't need
|
||||
if "num_batches_tracked" in mlx_key:
|
||||
continue
|
||||
|
||||
params[mlx_key] = mx.array(value)
|
||||
|
||||
return params
|
||||
|
||||
|
||||
def _params_to_nested(flat: dict[str, mx.array]) -> dict:
|
||||
"""Convert flat dot-separated keys to nested dict for mlx load_weights."""
|
||||
nested: dict = {}
|
||||
for key, value in flat.items():
|
||||
parts = key.split(".")
|
||||
current = nested
|
||||
for part in parts[:-1]:
|
||||
if part not in current:
|
||||
current[part] = {}
|
||||
current = current[part]
|
||||
current[parts[-1]] = value
|
||||
return nested
|
||||
|
||||
|
||||
def _build_and_load_decoder(
|
||||
output_dim: int,
|
||||
prefix: str,
|
||||
weights: dict[str, np.ndarray],
|
||||
) -> DecoderHead:
|
||||
"""Build an MLX DecoderHead and load converted weights."""
|
||||
decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim)
|
||||
flat_params = _load_decoder_weights(prefix, weights)
|
||||
nested_params = _params_to_nested(flat_params)
|
||||
decoder.load_weights(list(_flatten_nested(nested_params)))
|
||||
decoder.eval() # use running stats for BatchNorm (not batch stats)
|
||||
mx.eval(decoder.parameters()) # noqa: S307 — mx.eval is MLX lazy eval, not Python eval
|
||||
return decoder
|
||||
|
||||
|
||||
def _flatten_nested(d: dict, prefix: str = "") -> list[tuple[str, mx.array]]:
|
||||
"""Flatten nested dict to list of (dotted_key, array) pairs."""
|
||||
items: list[tuple[str, mx.array]] = []
|
||||
for k, v in d.items():
|
||||
full_key = f"{prefix}.{k}" if prefix else k
|
||||
if isinstance(v, dict):
|
||||
items.extend(_flatten_nested(v, full_key))
|
||||
else:
|
||||
items.append((full_key, v))
|
||||
return items
|
||||
|
||||
|
||||
def _run_decoder_parity(
|
||||
output_dim: int,
|
||||
prefix: str,
|
||||
expected_key: str,
|
||||
) -> None:
|
||||
"""Run a decoder parity test."""
|
||||
_skip_if_missing()
|
||||
|
||||
fixtures = dict(np.load(FIXTURE_PATH))
|
||||
weights = dict(np.load(WEIGHTS_PATH))
|
||||
|
||||
# Load features in NHWC
|
||||
features_nhwc = [mx.array(nchw_to_nhwc_np(fixtures[f"encoder_feature_{i}"])) for i in range(4)]
|
||||
|
||||
decoder = _build_and_load_decoder(output_dim, prefix, weights)
|
||||
result_nhwc = decoder(features_nhwc)
|
||||
mx.eval(result_nhwc) # noqa: S307 — mx.eval is MLX lazy eval, not Python eval
|
||||
|
||||
# Convert result back to NCHW for comparison
|
||||
result_nchw = nhwc_to_nchw_np(np.array(result_nhwc))
|
||||
expected = fixtures[expected_key]
|
||||
|
||||
max_abs_err = float(np.max(np.abs(result_nchw - expected)))
|
||||
mean_abs_err = float(np.mean(np.abs(result_nchw - expected)))
|
||||
print(f"\n{prefix} parity — max_abs: {max_abs_err:.6e}, mean_abs: {mean_abs_err:.6e}")
|
||||
|
||||
assert result_nchw.shape == expected.shape, (
|
||||
f"Shape mismatch: {result_nchw.shape} vs {expected.shape}"
|
||||
)
|
||||
assert max_abs_err < MAX_ABS_TOL, (
|
||||
f"Max abs error {max_abs_err:.6e} exceeds tolerance {MAX_ABS_TOL}"
|
||||
)
|
||||
|
||||
|
||||
def test_alpha_decoder_parity() -> None:
|
||||
"""MLX alpha decoder matches PyTorch within tolerance."""
|
||||
_run_decoder_parity(
|
||||
output_dim=1,
|
||||
prefix="alpha_decoder.",
|
||||
expected_key="alpha_logits",
|
||||
)
|
||||
|
||||
|
||||
def test_fg_decoder_parity() -> None:
|
||||
"""MLX foreground decoder matches PyTorch within tolerance."""
|
||||
_run_decoder_parity(
|
||||
output_dim=3,
|
||||
prefix="fg_decoder.",
|
||||
expected_key="fg_logits",
|
||||
)
|
||||
@ -1,94 +0,0 @@
|
||||
"""End-to-end parity tests: MLX GreenFormer vs PyTorch golden reference (Phase 5).
|
||||
|
||||
Loads golden.npz (PyTorch NCHW), runs same input through MLX model,
|
||||
compares all intermediate and final outputs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from corridorkey_mlx.model.corridorkey import GreenFormer
|
||||
from corridorkey_mlx.utils.layout import nchw_to_nhwc_np, nhwc_to_nchw_np
|
||||
|
||||
FIXTURE_PATH = Path("reference/fixtures/golden.npz")
|
||||
CHECKPOINT_PATH = Path("checkpoints/corridorkey_mlx.safetensors")
|
||||
IMG_SIZE = 512
|
||||
|
||||
# Tolerance tiers — coarse path inherits backbone drift, refiner adds more.
|
||||
# Backbone stages have up to ~0.01 max abs err (16 sequential blocks).
|
||||
# Final outputs pass through sigmoid which compresses errors.
|
||||
TOLERANCES: dict[str, float] = {
|
||||
"alpha_logits": 1e-3,
|
||||
"fg_logits": 1e-3,
|
||||
"alpha_logits_up": 1e-3,
|
||||
"fg_logits_up": 1e-3,
|
||||
"alpha_coarse": 1e-4,
|
||||
"fg_coarse": 1e-4,
|
||||
"delta_logits": 2e-3,
|
||||
"alpha_final": 2e-4,
|
||||
"fg_final": 2e-4,
|
||||
}
|
||||
|
||||
|
||||
def _skip_if_missing() -> None:
|
||||
if not FIXTURE_PATH.exists():
|
||||
pytest.skip("golden.npz not found — run dump_pytorch_reference.py first")
|
||||
if not CHECKPOINT_PATH.exists():
|
||||
pytest.skip("Checkpoint not found — run scripts/convert_weights.py first")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def model_outputs_and_fixtures() -> tuple[dict[str, mx.array], dict[str, np.ndarray]]:
|
||||
"""Load model, run forward pass with golden input, return (mlx_outputs, fixtures)."""
|
||||
_skip_if_missing()
|
||||
|
||||
fixtures = dict(np.load(FIXTURE_PATH))
|
||||
|
||||
# Load input: NCHW -> NHWC
|
||||
input_nhwc = mx.array(nchw_to_nhwc_np(fixtures["input"]))
|
||||
|
||||
model = GreenFormer(img_size=IMG_SIZE)
|
||||
model.load_checkpoint(CHECKPOINT_PATH)
|
||||
|
||||
outputs = model(input_nhwc)
|
||||
# materialize — mx.eval is MLX lazy graph evaluation, not Python eval
|
||||
mx.eval(outputs) # noqa: S307
|
||||
|
||||
return outputs, fixtures
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key", list(TOLERANCES.keys()))
|
||||
def test_parity(
|
||||
key: str,
|
||||
model_outputs_and_fixtures: tuple[dict[str, mx.array], dict[str, np.ndarray]],
|
||||
) -> None:
|
||||
"""MLX output matches PyTorch golden reference within tolerance."""
|
||||
outputs, fixtures = model_outputs_and_fixtures
|
||||
|
||||
assert key in outputs, f"Missing output key: {key}"
|
||||
assert key in fixtures, f"Missing fixture key: {key}"
|
||||
|
||||
# MLX output is NHWC, fixture is NCHW — convert MLX to NCHW for comparison
|
||||
mlx_nchw = nhwc_to_nchw_np(np.array(outputs[key]))
|
||||
expected_nchw = fixtures[key]
|
||||
|
||||
assert mlx_nchw.shape == expected_nchw.shape, (
|
||||
f"{key} shape mismatch: MLX {mlx_nchw.shape} vs PyTorch {expected_nchw.shape}"
|
||||
)
|
||||
|
||||
abs_err = np.abs(mlx_nchw - expected_nchw)
|
||||
max_abs = float(np.max(abs_err))
|
||||
mean_abs = float(np.mean(abs_err))
|
||||
tol = TOLERANCES[key]
|
||||
|
||||
print(
|
||||
f"\n{key:<20s} | shape={mlx_nchw.shape} | "
|
||||
f"max_abs={max_abs:.6e} | mean_abs={mean_abs:.6e} | tol={tol:.1e}"
|
||||
)
|
||||
|
||||
assert max_abs < tol, f"{key}: max abs error {max_abs:.6e} exceeds tolerance {tol:.1e}"
|
||||
@ -1,66 +0,0 @@
|
||||
"""End-to-end smoke test for GreenFormer (Phase 5).
|
||||
|
||||
Verifies the full pipeline works with random weights — no checkpoint needed.
|
||||
Tests model construction, forward pass, and basic output sanity.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import mlx.core as mx
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from corridorkey_mlx.io.image import (
|
||||
postprocess_alpha,
|
||||
postprocess_foreground,
|
||||
preprocess,
|
||||
)
|
||||
from corridorkey_mlx.model.corridorkey import GreenFormer
|
||||
|
||||
IMG_SIZE = 256 # smaller for speed
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def model() -> GreenFormer:
|
||||
return GreenFormer(img_size=IMG_SIZE)
|
||||
|
||||
|
||||
def test_forward_with_preprocessed_input(model: GreenFormer) -> None:
|
||||
"""Full pipeline: numpy image -> preprocess -> model -> postprocess."""
|
||||
rgb = np.random.rand(IMG_SIZE, IMG_SIZE, 3).astype(np.float32)
|
||||
alpha_hint = np.random.rand(IMG_SIZE, IMG_SIZE, 1).astype(np.float32)
|
||||
|
||||
x = preprocess(rgb, alpha_hint)
|
||||
assert x.shape == (1, IMG_SIZE, IMG_SIZE, 4)
|
||||
|
||||
out = model(x)
|
||||
# materialize — mx.eval is MLX's lazy graph evaluation, not Python eval
|
||||
mx.eval(out)
|
||||
|
||||
alpha = postprocess_alpha(out["alpha_final"])
|
||||
fg = postprocess_foreground(out["fg_final"])
|
||||
|
||||
assert alpha.shape == (IMG_SIZE, IMG_SIZE)
|
||||
assert alpha.dtype == np.uint8
|
||||
assert fg.shape == (IMG_SIZE, IMG_SIZE, 3)
|
||||
assert fg.dtype == np.uint8
|
||||
|
||||
|
||||
def test_deterministic_output(model: GreenFormer) -> None:
|
||||
"""Same input produces same output."""
|
||||
mx.random.seed(123)
|
||||
x = mx.random.normal((1, IMG_SIZE, IMG_SIZE, 4))
|
||||
# materialize — mx.eval is MLX's lazy graph evaluation, not Python eval
|
||||
mx.eval(x)
|
||||
|
||||
out1 = model(x)
|
||||
mx.eval(out1)
|
||||
out2 = model(x)
|
||||
mx.eval(out2)
|
||||
|
||||
for key in out1:
|
||||
np.testing.assert_array_equal(
|
||||
np.array(out1[key]),
|
||||
np.array(out2[key]),
|
||||
err_msg=f"{key} not deterministic",
|
||||
)
|
||||
@ -2,8 +2,6 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
@ -13,9 +11,7 @@ from corridorkey_mlx.engine import (
|
||||
_validate_mask,
|
||||
)
|
||||
|
||||
CHECKPOINT = Path("checkpoints/corridorkey_mlx.safetensors")
|
||||
HAS_CHECKPOINT = CHECKPOINT.exists()
|
||||
|
||||
from .conftest import MLX_CHECKPOINT_PATH, has_checkpoint
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Input validation (no checkpoint needed)
|
||||
@ -87,14 +83,14 @@ class TestEngineInit:
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_CHECKPOINT, reason="Checkpoint not available")
|
||||
@has_checkpoint
|
||||
class TestEngineIntegration:
|
||||
"""Integration tests that load the real model."""
|
||||
|
||||
@pytest.fixture(scope="class")
|
||||
def engine(self) -> CorridorKeyMLXEngine:
|
||||
return CorridorKeyMLXEngine(
|
||||
checkpoint_path=CHECKPOINT,
|
||||
checkpoint_path=MLX_CHECKPOINT_PATH,
|
||||
img_size=512,
|
||||
compile=False,
|
||||
)
|
||||
@ -128,10 +124,42 @@ class TestEngineIntegration:
|
||||
result = engine.process_frame(image, mask)
|
||||
assert "alpha" in result
|
||||
|
||||
def test_refiner_scale_zero_returns_coarse(
|
||||
self, engine: CorridorKeyMLXEngine
|
||||
) -> None:
|
||||
def test_refiner_scale_zero_returns_coarse(self, engine: CorridorKeyMLXEngine) -> None:
|
||||
image = np.random.default_rng(42).integers(0, 256, (64, 64, 3), dtype=np.uint8)
|
||||
mask = np.random.default_rng(42).integers(0, 256, (64, 64), dtype=np.uint8)
|
||||
result = engine.process_frame(image, mask, refiner_scale=0.0)
|
||||
assert result["alpha"].shape == (64, 64)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Full 2048 inference (slow, requires checkpoint)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@has_checkpoint
|
||||
def test_smoke_2048_full() -> None:
|
||||
"""Full 2048 inference with real checkpoint."""
|
||||
engine = CorridorKeyMLXEngine(
|
||||
checkpoint_path=MLX_CHECKPOINT_PATH,
|
||||
img_size=2048,
|
||||
compile=False,
|
||||
)
|
||||
|
||||
rng = np.random.default_rng(42)
|
||||
rgb = rng.integers(0, 256, (2048, 2048, 3), dtype=np.uint8)
|
||||
mask = rng.integers(0, 256, (2048, 2048), dtype=np.uint8)
|
||||
|
||||
result = engine.process_frame(rgb, mask)
|
||||
|
||||
assert result["alpha"].shape == (2048, 2048)
|
||||
assert result["alpha"].dtype == np.uint8
|
||||
assert result["fg"].shape == (2048, 2048, 3)
|
||||
assert result["comp"].shape == (2048, 2048, 3)
|
||||
|
||||
for key in ("alpha", "fg", "comp"):
|
||||
arr = result[key]
|
||||
assert not np.isnan(arr).any(), f"{key} contains NaN"
|
||||
assert not np.isinf(arr).any(), f"{key} contains Inf"
|
||||
|
||||
assert result["alpha"].min() != result["alpha"].max(), "alpha is constant"
|
||||
|
||||
@ -1,77 +0,0 @@
|
||||
"""Forward pass shape and smoke tests for GreenFormer (Phase 5).
|
||||
|
||||
Uses random weights — no checkpoint needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import mlx.core as mx
|
||||
import pytest
|
||||
|
||||
from corridorkey_mlx.model.corridorkey import GreenFormer
|
||||
|
||||
IMG_SIZE = 512
|
||||
BATCH = 1
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def model_and_output() -> tuple[GreenFormer, dict[str, mx.array]]:
|
||||
"""Build model with random weights and run forward pass once."""
|
||||
model = GreenFormer(img_size=IMG_SIZE)
|
||||
x = mx.random.normal((BATCH, IMG_SIZE, IMG_SIZE, 4))
|
||||
out = model(x)
|
||||
# materialize all outputs — mx.eval is MLX lazy graph evaluation, not Python eval
|
||||
mx.eval(out)
|
||||
return model, out
|
||||
|
||||
|
||||
EXPECTED_SHAPES: dict[str, tuple[int, ...]] = {
|
||||
"alpha_logits": (BATCH, IMG_SIZE // 4, IMG_SIZE // 4, 1),
|
||||
"fg_logits": (BATCH, IMG_SIZE // 4, IMG_SIZE // 4, 3),
|
||||
"alpha_logits_up": (BATCH, IMG_SIZE, IMG_SIZE, 1),
|
||||
"fg_logits_up": (BATCH, IMG_SIZE, IMG_SIZE, 3),
|
||||
"alpha_coarse": (BATCH, IMG_SIZE, IMG_SIZE, 1),
|
||||
"fg_coarse": (BATCH, IMG_SIZE, IMG_SIZE, 3),
|
||||
"delta_logits": (BATCH, IMG_SIZE, IMG_SIZE, 4),
|
||||
"alpha_final": (BATCH, IMG_SIZE, IMG_SIZE, 1),
|
||||
"fg_final": (BATCH, IMG_SIZE, IMG_SIZE, 3),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key,expected_shape", EXPECTED_SHAPES.items())
|
||||
def test_output_shapes(
|
||||
key: str,
|
||||
expected_shape: tuple[int, ...],
|
||||
model_and_output: tuple[GreenFormer, dict[str, mx.array]],
|
||||
) -> None:
|
||||
_, out = model_and_output
|
||||
assert key in out, f"Missing output key: {key}"
|
||||
assert out[key].shape == expected_shape, (
|
||||
f"{key}: expected {expected_shape}, got {out[key].shape}"
|
||||
)
|
||||
|
||||
|
||||
def test_all_keys_present(
|
||||
model_and_output: tuple[GreenFormer, dict[str, mx.array]],
|
||||
) -> None:
|
||||
_, out = model_and_output
|
||||
assert set(out.keys()) == set(EXPECTED_SHAPES.keys())
|
||||
|
||||
|
||||
def test_coarse_probs_in_range(
|
||||
model_and_output: tuple[GreenFormer, dict[str, mx.array]],
|
||||
) -> None:
|
||||
"""Sigmoid outputs must be in [0, 1]."""
|
||||
_, out = model_and_output
|
||||
for key in ("alpha_coarse", "fg_coarse", "alpha_final", "fg_final"):
|
||||
arr = out[key]
|
||||
assert float(mx.min(arr)) >= 0.0, f"{key} has values < 0"
|
||||
assert float(mx.max(arr)) <= 1.0, f"{key} has values > 1"
|
||||
|
||||
|
||||
def test_output_dtype(
|
||||
model_and_output: tuple[GreenFormer, dict[str, mx.array]],
|
||||
) -> None:
|
||||
_, out = model_and_output
|
||||
for key, arr in out.items():
|
||||
assert arr.dtype == mx.float32, f"{key}: expected float32, got {arr.dtype}"
|
||||
@ -1,79 +0,0 @@
|
||||
"""Parity tests: MLX Hiera backbone vs PyTorch reference (Phase 4).
|
||||
|
||||
Loads golden input + encoder features from fixtures, runs MLX backbone
|
||||
with checkpoint weights, compares each stage output.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from corridorkey_mlx.model.hiera import HieraBackbone
|
||||
from corridorkey_mlx.utils.layout import nchw_to_nhwc_np, nhwc_to_nchw_np
|
||||
|
||||
FIXTURE_PATH = Path("reference/fixtures/golden.npz")
|
||||
CHECKPOINT_PATH = Path("checkpoints/corridorkey_mlx.safetensors")
|
||||
IMG_SIZE = 512
|
||||
NUM_STAGES = 4
|
||||
|
||||
# Stage 2 runs 16 consecutive blocks — float32 drift accumulates on Metal vs CPU.
|
||||
# Mean error stays < 1e-4; max outliers can reach ~0.01 in the deepest stage.
|
||||
MAX_ABS_TOL = 2e-2
|
||||
|
||||
|
||||
def _skip_if_missing() -> None:
|
||||
if not FIXTURE_PATH.exists():
|
||||
pytest.skip("Fixture files not found — run dump_pytorch_reference.py first")
|
||||
if not CHECKPOINT_PATH.exists():
|
||||
pytest.skip("Checkpoint not found — run scripts/convert_weights.py first")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def backbone_and_fixtures() -> tuple[list[mx.array], dict[str, np.ndarray]]:
|
||||
"""Load backbone once, return (mlx_features, fixtures)."""
|
||||
_skip_if_missing()
|
||||
|
||||
fixtures = dict(np.load(FIXTURE_PATH))
|
||||
|
||||
# Load input: NCHW -> NHWC
|
||||
input_nchw = fixtures["input"]
|
||||
input_nhwc = mx.array(nchw_to_nhwc_np(input_nchw))
|
||||
|
||||
backbone = HieraBackbone(img_size=IMG_SIZE)
|
||||
backbone.load_checkpoint(CHECKPOINT_PATH)
|
||||
|
||||
features = backbone(input_nhwc)
|
||||
# materialize all features — mx.eval is MLX lazy evaluation, not Python eval
|
||||
mx.eval(features) # noqa: S307
|
||||
|
||||
return features, fixtures
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stage_idx", range(NUM_STAGES))
|
||||
def test_stage_parity(
|
||||
stage_idx: int,
|
||||
backbone_and_fixtures: tuple[list[mx.array], dict[str, np.ndarray]],
|
||||
) -> None:
|
||||
"""MLX backbone stage output matches PyTorch within tolerance."""
|
||||
features, fixtures = backbone_and_fixtures
|
||||
|
||||
expected_nchw = fixtures[f"encoder_feature_{stage_idx}"]
|
||||
result_nhwc = features[stage_idx]
|
||||
result_nchw = nhwc_to_nchw_np(np.array(result_nhwc))
|
||||
|
||||
assert result_nchw.shape == expected_nchw.shape, (
|
||||
f"Stage {stage_idx} shape mismatch: {result_nchw.shape} vs {expected_nchw.shape}"
|
||||
)
|
||||
|
||||
abs_err = np.abs(result_nchw - expected_nchw)
|
||||
max_abs_err = float(np.max(abs_err))
|
||||
mean_abs_err = float(np.mean(abs_err))
|
||||
print(f"\nStage {stage_idx} parity — max_abs: {max_abs_err:.6e}, mean_abs: {mean_abs_err:.6e}")
|
||||
|
||||
assert max_abs_err < MAX_ABS_TOL, (
|
||||
f"Stage {stage_idx} max abs error {max_abs_err:.6e} exceeds tolerance {MAX_ABS_TOL}"
|
||||
)
|
||||
@ -1,46 +0,0 @@
|
||||
"""Shape contract tests for Hiera backbone (Phase 4).
|
||||
|
||||
No checkpoint needed — verifies structural correctness with random weights.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import mlx.core as mx
|
||||
import pytest
|
||||
|
||||
from corridorkey_mlx.model.hiera import HieraBackbone, HieraPatchEmbed
|
||||
|
||||
IMG_SIZE = 512
|
||||
NUM_FEATURES = 4
|
||||
EXPECTED_SHAPES = [
|
||||
(1, 128, 128, 112), # stride 4
|
||||
(1, 64, 64, 224), # stride 8
|
||||
(1, 32, 32, 448), # stride 16
|
||||
(1, 16, 16, 896), # stride 32
|
||||
]
|
||||
|
||||
|
||||
def test_patch_embed_output_shape() -> None:
|
||||
"""PatchEmbed produces (B, N, C) with N = (H/4)*(W/4)."""
|
||||
patch_embed = HieraPatchEmbed()
|
||||
x = mx.zeros((1, IMG_SIZE, IMG_SIZE, 4))
|
||||
out = patch_embed(x)
|
||||
expected_n = (IMG_SIZE // 4) * (IMG_SIZE // 4)
|
||||
assert out.shape == (1, expected_n, 112)
|
||||
|
||||
|
||||
def test_backbone_returns_four_features() -> None:
|
||||
"""Backbone returns exactly 4 feature maps."""
|
||||
backbone = HieraBackbone(img_size=IMG_SIZE)
|
||||
x = mx.zeros((1, IMG_SIZE, IMG_SIZE, 4))
|
||||
features = backbone(x)
|
||||
assert len(features) == NUM_FEATURES
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stage_idx", range(NUM_FEATURES))
|
||||
def test_feature_map_shape(stage_idx: int) -> None:
|
||||
"""Each stage feature map has correct (B, H, W, C) shape."""
|
||||
backbone = HieraBackbone(img_size=IMG_SIZE)
|
||||
x = mx.zeros((1, IMG_SIZE, IMG_SIZE, 4))
|
||||
features = backbone(x)
|
||||
assert features[stage_idx].shape == EXPECTED_SHAPES[stage_idx]
|
||||
@ -1,12 +0,0 @@
|
||||
"""Smoke test: verify package imports succeed."""
|
||||
|
||||
|
||||
def test_package_imports() -> None:
|
||||
import corridorkey_mlx
|
||||
import corridorkey_mlx.convert
|
||||
import corridorkey_mlx.inference
|
||||
import corridorkey_mlx.io
|
||||
import corridorkey_mlx.model
|
||||
import corridorkey_mlx.utils
|
||||
|
||||
assert corridorkey_mlx.__version__ == "0.1.0"
|
||||
194
tests/test_model_contract.py
Normal file
194
tests/test_model_contract.py
Normal file
@ -0,0 +1,194 @@
|
||||
"""Model contract tests: imports, shapes, outputs, determinism.
|
||||
|
||||
Consolidated from: test_import.py, test_greenformer_forward.py,
|
||||
test_hiera_stage_shapes.py, test_end_to_end_smoke.py, test_smoke_2048.py (wiring).
|
||||
|
||||
No checkpoint needed — uses random weights only.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import mlx.core as mx
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from corridorkey_mlx.model.corridorkey import GreenFormer
|
||||
from corridorkey_mlx.model.hiera import HieraBackbone, HieraPatchEmbed
|
||||
|
||||
from .conftest import IMG_SIZE, SMALL_IMG_SIZE
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Package imports
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_package_imports() -> None:
|
||||
import corridorkey_mlx
|
||||
import corridorkey_mlx.convert # noqa: F401
|
||||
import corridorkey_mlx.inference # noqa: F401
|
||||
import corridorkey_mlx.io # noqa: F401
|
||||
import corridorkey_mlx.model # noqa: F401
|
||||
import corridorkey_mlx.utils # noqa: F401
|
||||
|
||||
assert corridorkey_mlx.__version__ == "0.1.0"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GreenFormer output contract (random weights)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
EXPECTED_SHAPES: dict[str, tuple[int, ...]] = {
|
||||
"alpha_logits": (1, IMG_SIZE // 4, IMG_SIZE // 4, 1),
|
||||
"fg_logits": (1, IMG_SIZE // 4, IMG_SIZE // 4, 3),
|
||||
"alpha_logits_up": (1, IMG_SIZE, IMG_SIZE, 1),
|
||||
"fg_logits_up": (1, IMG_SIZE, IMG_SIZE, 3),
|
||||
"alpha_coarse": (1, IMG_SIZE, IMG_SIZE, 1),
|
||||
"fg_coarse": (1, IMG_SIZE, IMG_SIZE, 3),
|
||||
"delta_logits": (1, IMG_SIZE, IMG_SIZE, 4),
|
||||
"alpha_final": (1, IMG_SIZE, IMG_SIZE, 1),
|
||||
"fg_final": (1, IMG_SIZE, IMG_SIZE, 3),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def model_output() -> dict[str, mx.array]:
|
||||
"""Forward pass with random weights, materialized once."""
|
||||
model = GreenFormer(img_size=IMG_SIZE)
|
||||
x = mx.random.normal((1, IMG_SIZE, IMG_SIZE, 4))
|
||||
out = model(x)
|
||||
# mx.eval is MLX array materialization, not Python eval()
|
||||
mx.eval(out) # noqa: S307
|
||||
return out
|
||||
|
||||
|
||||
def test_all_keys_present(model_output: dict[str, mx.array]) -> None:
|
||||
assert set(model_output.keys()) == set(EXPECTED_SHAPES.keys())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key,expected_shape", EXPECTED_SHAPES.items())
|
||||
def test_output_shapes(
|
||||
key: str,
|
||||
expected_shape: tuple[int, ...],
|
||||
model_output: dict[str, mx.array],
|
||||
) -> None:
|
||||
assert model_output[key].shape == expected_shape, (
|
||||
f"{key}: expected {expected_shape}, got {model_output[key].shape}"
|
||||
)
|
||||
|
||||
|
||||
def test_output_dtype_float32(model_output: dict[str, mx.array]) -> None:
|
||||
for key, arr in model_output.items():
|
||||
assert arr.dtype == mx.float32, f"{key}: expected float32, got {arr.dtype}"
|
||||
|
||||
|
||||
def test_sigmoid_outputs_in_range(model_output: dict[str, mx.array]) -> None:
|
||||
"""Post-sigmoid outputs must be in [0, 1]."""
|
||||
for key in ("alpha_coarse", "fg_coarse", "alpha_final", "fg_final"):
|
||||
arr = model_output[key]
|
||||
assert float(mx.min(arr)) >= 0.0, f"{key} has values < 0"
|
||||
assert float(mx.max(arr)) <= 1.0, f"{key} has values > 1"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Backbone shape contract
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
BACKBONE_SHAPES = [
|
||||
(1, 128, 128, 112), # stride 4
|
||||
(1, 64, 64, 224), # stride 8
|
||||
(1, 32, 32, 448), # stride 16
|
||||
(1, 16, 16, 896), # stride 32
|
||||
]
|
||||
|
||||
|
||||
def test_patch_embed_output_shape() -> None:
|
||||
patch_embed = HieraPatchEmbed()
|
||||
x = mx.zeros((1, IMG_SIZE, IMG_SIZE, 4))
|
||||
out = patch_embed(x)
|
||||
expected_n = (IMG_SIZE // 4) * (IMG_SIZE // 4)
|
||||
assert out.shape == (1, expected_n, 112)
|
||||
|
||||
|
||||
def test_backbone_returns_four_features() -> None:
|
||||
backbone = HieraBackbone(img_size=IMG_SIZE)
|
||||
x = mx.zeros((1, IMG_SIZE, IMG_SIZE, 4))
|
||||
features = backbone(x)
|
||||
assert len(features) == 4
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stage_idx", range(4))
|
||||
def test_backbone_feature_shapes(stage_idx: int) -> None:
|
||||
backbone = HieraBackbone(img_size=IMG_SIZE)
|
||||
x = mx.zeros((1, IMG_SIZE, IMG_SIZE, 4))
|
||||
features = backbone(x)
|
||||
assert features[stage_idx].shape == BACKBONE_SHAPES[stage_idx]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Full pipeline roundtrip (preprocess -> model -> postprocess)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_pipeline_roundtrip() -> None:
|
||||
"""numpy -> preprocess -> model -> postprocess -> uint8."""
|
||||
from corridorkey_mlx.io.image import postprocess_alpha, postprocess_foreground, preprocess
|
||||
|
||||
model = GreenFormer(img_size=SMALL_IMG_SIZE)
|
||||
rgb = np.random.rand(SMALL_IMG_SIZE, SMALL_IMG_SIZE, 3).astype(np.float32)
|
||||
alpha_hint = np.random.rand(SMALL_IMG_SIZE, SMALL_IMG_SIZE, 1).astype(np.float32)
|
||||
|
||||
x = preprocess(rgb, alpha_hint)
|
||||
assert x.shape == (1, SMALL_IMG_SIZE, SMALL_IMG_SIZE, 4)
|
||||
|
||||
out = model(x)
|
||||
mx.eval(out) # noqa: S307
|
||||
|
||||
alpha = postprocess_alpha(out["alpha_final"])
|
||||
fg = postprocess_foreground(out["fg_final"])
|
||||
|
||||
assert alpha.shape == (SMALL_IMG_SIZE, SMALL_IMG_SIZE)
|
||||
assert alpha.dtype == np.uint8
|
||||
assert fg.shape == (SMALL_IMG_SIZE, SMALL_IMG_SIZE, 3)
|
||||
assert fg.dtype == np.uint8
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Determinism
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_deterministic_output() -> None:
|
||||
"""Same input -> same output."""
|
||||
model = GreenFormer(img_size=SMALL_IMG_SIZE)
|
||||
mx.random.seed(123)
|
||||
x = mx.random.normal((1, SMALL_IMG_SIZE, SMALL_IMG_SIZE, 4))
|
||||
mx.eval(x) # noqa: S307
|
||||
|
||||
out1 = model(x)
|
||||
mx.eval(out1) # noqa: S307
|
||||
out2 = model(x)
|
||||
mx.eval(out2) # noqa: S307
|
||||
|
||||
for key in out1:
|
||||
np.testing.assert_array_equal(
|
||||
np.array(out1[key]), np.array(out2[key]), err_msg=f"{key} not deterministic"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# No NaN/Inf (wiring check)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_no_nan_inf() -> None:
|
||||
"""Random-weight forward produces finite outputs."""
|
||||
model = GreenFormer(img_size=SMALL_IMG_SIZE)
|
||||
x = mx.random.normal((1, SMALL_IMG_SIZE, SMALL_IMG_SIZE, 4))
|
||||
mx.eval(x) # noqa: S307
|
||||
out = model(x)
|
||||
mx.eval(out) # noqa: S307
|
||||
|
||||
for key in ("alpha_final", "fg_final"):
|
||||
arr = np.array(out[key])
|
||||
assert not np.isnan(arr).any(), f"{key} contains NaN"
|
||||
assert not np.isinf(arr).any(), f"{key} contains Inf"
|
||||
140
tests/test_parity.py
Normal file
140
tests/test_parity.py
Normal file
@ -0,0 +1,140 @@
|
||||
"""Consolidated golden reference parity tests.
|
||||
|
||||
Merged from: test_end_to_end_parity.py, test_decoder_parity.py,
|
||||
test_refiner_parity.py, test_hiera_stage_parity.py, test_reference_fixtures.py.
|
||||
|
||||
Loads full model with checkpoint, runs forward on golden input,
|
||||
compares all outputs against PyTorch golden.npz references.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import mlx.core as mx
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from corridorkey_mlx.model.corridorkey import GreenFormer
|
||||
from corridorkey_mlx.utils.layout import nchw_to_nhwc_np, nhwc_to_nchw_np
|
||||
|
||||
from .conftest import (
|
||||
GOLDEN_PATH,
|
||||
IMG_SIZE,
|
||||
MLX_CHECKPOINT_PATH,
|
||||
PARITY_TOL_BACKBONE,
|
||||
)
|
||||
|
||||
NUM_BACKBONE_STAGES = 4
|
||||
|
||||
# Per-output tolerances — backbone drift cascades through decoder/refiner.
|
||||
OUTPUT_TOLERANCES: dict[str, float] = {
|
||||
"alpha_logits": 1e-3,
|
||||
"fg_logits": 1e-3,
|
||||
"alpha_logits_up": 1e-3,
|
||||
"fg_logits_up": 1e-3,
|
||||
"alpha_coarse": 1e-4,
|
||||
"fg_coarse": 1e-4,
|
||||
"delta_logits": 2e-3,
|
||||
"alpha_final": 2e-4,
|
||||
"fg_final": 2e-4,
|
||||
}
|
||||
|
||||
|
||||
def _skip_if_missing() -> None:
|
||||
if not GOLDEN_PATH.exists():
|
||||
pytest.skip("golden.npz not found")
|
||||
if not MLX_CHECKPOINT_PATH.exists():
|
||||
pytest.skip("MLX checkpoint not found")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def model_outputs_and_fixtures() -> tuple[dict[str, mx.array], dict[str, np.ndarray]]:
|
||||
"""Load model, run forward with golden input, return (mlx_outputs, fixtures)."""
|
||||
_skip_if_missing()
|
||||
|
||||
fixtures = dict(np.load(GOLDEN_PATH))
|
||||
|
||||
# NCHW -> NHWC for MLX
|
||||
input_nhwc = mx.array(nchw_to_nhwc_np(fixtures["input"]))
|
||||
|
||||
model = GreenFormer(img_size=IMG_SIZE)
|
||||
model.load_checkpoint(MLX_CHECKPOINT_PATH)
|
||||
|
||||
outputs = model(input_nhwc)
|
||||
# mx.eval is MLX array materialization, not Python eval()
|
||||
mx.eval(outputs) # noqa: S307
|
||||
|
||||
return outputs, fixtures
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Backbone stage parity (4 stages)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def backbone_features() -> tuple[list[mx.array], dict[str, np.ndarray]]:
|
||||
"""Run backbone alone to get intermediate stage features."""
|
||||
_skip_if_missing()
|
||||
|
||||
fixtures = dict(np.load(GOLDEN_PATH))
|
||||
input_nhwc = mx.array(nchw_to_nhwc_np(fixtures["input"]))
|
||||
|
||||
from corridorkey_mlx.model.hiera import HieraBackbone
|
||||
|
||||
backbone = HieraBackbone(img_size=IMG_SIZE)
|
||||
backbone.load_checkpoint(MLX_CHECKPOINT_PATH)
|
||||
|
||||
features = backbone(input_nhwc)
|
||||
mx.eval(features) # noqa: S307
|
||||
|
||||
return features, fixtures
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stage_idx", range(NUM_BACKBONE_STAGES))
|
||||
def test_backbone_stage_parity(
|
||||
stage_idx: int,
|
||||
backbone_features: tuple[list[mx.array], dict[str, np.ndarray]],
|
||||
) -> None:
|
||||
"""MLX backbone stage output matches PyTorch within tolerance."""
|
||||
features, fixtures = backbone_features
|
||||
|
||||
expected_nchw = fixtures[f"encoder_feature_{stage_idx}"]
|
||||
result_nchw = nhwc_to_nchw_np(np.array(features[stage_idx]))
|
||||
|
||||
assert result_nchw.shape == expected_nchw.shape, (
|
||||
f"Stage {stage_idx} shape mismatch: {result_nchw.shape} vs {expected_nchw.shape}"
|
||||
)
|
||||
|
||||
max_abs_err = float(np.max(np.abs(result_nchw - expected_nchw)))
|
||||
assert max_abs_err < PARITY_TOL_BACKBONE, (
|
||||
f"Stage {stage_idx} max_abs={max_abs_err:.2e} > {PARITY_TOL_BACKBONE}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# E2E output parity (decoder + refiner outputs)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key", list(OUTPUT_TOLERANCES.keys()))
|
||||
def test_output_parity(
|
||||
key: str,
|
||||
model_outputs_and_fixtures: tuple[dict[str, mx.array], dict[str, np.ndarray]],
|
||||
) -> None:
|
||||
"""MLX output matches PyTorch golden reference within tolerance."""
|
||||
outputs, fixtures = model_outputs_and_fixtures
|
||||
|
||||
assert key in outputs, f"Missing MLX output key: {key}"
|
||||
assert key in fixtures, f"Missing fixture key: {key}"
|
||||
|
||||
mlx_nchw = nhwc_to_nchw_np(np.array(outputs[key]))
|
||||
expected_nchw = fixtures[key]
|
||||
|
||||
assert mlx_nchw.shape == expected_nchw.shape, (
|
||||
f"{key} shape mismatch: {mlx_nchw.shape} vs {expected_nchw.shape}"
|
||||
)
|
||||
|
||||
max_abs = float(np.max(np.abs(mlx_nchw - expected_nchw)))
|
||||
tol = OUTPUT_TOLERANCES[key]
|
||||
|
||||
assert max_abs < tol, f"{key}: max_abs={max_abs:.2e} > tol={tol:.1e}"
|
||||
@ -1,130 +0,0 @@
|
||||
"""Parity tests: PyTorch reference fixtures (Phase 1).
|
||||
|
||||
These tests validate that dump_pytorch_reference.py produces
|
||||
fixtures with expected shapes and dtypes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
FIXTURE_PATH = Path("reference/fixtures/golden.npz")
|
||||
|
||||
IMG_SIZE = 512
|
||||
BATCH = 1
|
||||
|
||||
# Expected shapes for each tensor in the fixture file
|
||||
EXPECTED_SHAPES: dict[str, tuple[int, ...]] = {
|
||||
"input": (BATCH, 4, IMG_SIZE, IMG_SIZE),
|
||||
# Backbone features at strides 4, 8, 16, 32
|
||||
"encoder_feature_0": (BATCH, 112, IMG_SIZE // 4, IMG_SIZE // 4),
|
||||
"encoder_feature_1": (BATCH, 224, IMG_SIZE // 8, IMG_SIZE // 8),
|
||||
"encoder_feature_2": (BATCH, 448, IMG_SIZE // 16, IMG_SIZE // 16),
|
||||
"encoder_feature_3": (BATCH, 896, IMG_SIZE // 32, IMG_SIZE // 32),
|
||||
# Decoder outputs at H/4 resolution
|
||||
"alpha_logits": (BATCH, 1, IMG_SIZE // 4, IMG_SIZE // 4),
|
||||
"fg_logits": (BATCH, 3, IMG_SIZE // 4, IMG_SIZE // 4),
|
||||
# Upsampled to full resolution
|
||||
"alpha_logits_up": (BATCH, 1, IMG_SIZE, IMG_SIZE),
|
||||
"fg_logits_up": (BATCH, 3, IMG_SIZE, IMG_SIZE),
|
||||
# Coarse predictions (after sigmoid)
|
||||
"alpha_coarse": (BATCH, 1, IMG_SIZE, IMG_SIZE),
|
||||
"fg_coarse": (BATCH, 3, IMG_SIZE, IMG_SIZE),
|
||||
# Refiner
|
||||
"delta_logits": (BATCH, 4, IMG_SIZE, IMG_SIZE),
|
||||
# Final outputs
|
||||
"alpha_final": (BATCH, 1, IMG_SIZE, IMG_SIZE),
|
||||
"fg_final": (BATCH, 3, IMG_SIZE, IMG_SIZE),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def fixtures() -> dict[str, np.ndarray]:
|
||||
if not FIXTURE_PATH.exists():
|
||||
pytest.skip(f"Fixture file not found: {FIXTURE_PATH}")
|
||||
return dict(np.load(FIXTURE_PATH))
|
||||
|
||||
|
||||
class TestFixtureCompleteness:
|
||||
"""All expected tensors exist in the fixture file."""
|
||||
|
||||
def test_all_keys_present(self, fixtures: dict[str, np.ndarray]) -> None:
|
||||
missing = set(EXPECTED_SHAPES.keys()) - set(fixtures.keys())
|
||||
assert not missing, f"Missing fixture keys: {missing}"
|
||||
|
||||
def test_no_extra_keys(self, fixtures: dict[str, np.ndarray]) -> None:
|
||||
extra = set(fixtures.keys()) - set(EXPECTED_SHAPES.keys())
|
||||
assert not extra, f"Unexpected fixture keys: {extra}"
|
||||
|
||||
|
||||
class TestBackboneFeaturesShape:
|
||||
"""4 feature maps from Hiera backbone have expected shapes."""
|
||||
|
||||
@pytest.mark.parametrize("idx", range(4))
|
||||
def test_feature_shape(self, fixtures: dict[str, np.ndarray], idx: int) -> None:
|
||||
key = f"encoder_feature_{idx}"
|
||||
assert fixtures[key].shape == EXPECTED_SHAPES[key]
|
||||
|
||||
@pytest.mark.parametrize("idx", range(4))
|
||||
def test_feature_dtype(self, fixtures: dict[str, np.ndarray], idx: int) -> None:
|
||||
key = f"encoder_feature_{idx}"
|
||||
assert fixtures[key].dtype == np.float32
|
||||
|
||||
|
||||
class TestCoarsePredictionsShape:
|
||||
"""Alpha (1ch) and foreground (3ch) coarse logits/probs exist with correct shapes."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"key",
|
||||
[
|
||||
"alpha_logits",
|
||||
"fg_logits",
|
||||
"alpha_logits_up",
|
||||
"fg_logits_up",
|
||||
"alpha_coarse",
|
||||
"fg_coarse",
|
||||
],
|
||||
)
|
||||
def test_shape(self, fixtures: dict[str, np.ndarray], key: str) -> None:
|
||||
assert fixtures[key].shape == EXPECTED_SHAPES[key]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"key",
|
||||
[
|
||||
"alpha_logits",
|
||||
"fg_logits",
|
||||
"alpha_logits_up",
|
||||
"fg_logits_up",
|
||||
"alpha_coarse",
|
||||
"fg_coarse",
|
||||
],
|
||||
)
|
||||
def test_dtype(self, fixtures: dict[str, np.ndarray], key: str) -> None:
|
||||
assert fixtures[key].dtype == np.float32
|
||||
|
||||
@pytest.mark.parametrize("key", ["alpha_coarse", "fg_coarse"])
|
||||
def test_sigmoid_range(self, fixtures: dict[str, np.ndarray], key: str) -> None:
|
||||
"""Coarse predictions (post-sigmoid) should be in [0, 1]."""
|
||||
assert fixtures[key].min() >= 0.0
|
||||
assert fixtures[key].max() <= 1.0
|
||||
|
||||
|
||||
class TestRefinerOutputsShape:
|
||||
"""Delta logits and final alpha/fg have expected shapes."""
|
||||
|
||||
@pytest.mark.parametrize("key", ["delta_logits", "alpha_final", "fg_final"])
|
||||
def test_shape(self, fixtures: dict[str, np.ndarray], key: str) -> None:
|
||||
assert fixtures[key].shape == EXPECTED_SHAPES[key]
|
||||
|
||||
@pytest.mark.parametrize("key", ["delta_logits", "alpha_final", "fg_final"])
|
||||
def test_dtype(self, fixtures: dict[str, np.ndarray], key: str) -> None:
|
||||
assert fixtures[key].dtype == np.float32
|
||||
|
||||
@pytest.mark.parametrize("key", ["alpha_final", "fg_final"])
|
||||
def test_final_sigmoid_range(self, fixtures: dict[str, np.ndarray], key: str) -> None:
|
||||
"""Final predictions (post-sigmoid) should be in [0, 1]."""
|
||||
assert fixtures[key].min() >= 0.0
|
||||
assert fixtures[key].max() <= 1.0
|
||||
@ -1,154 +0,0 @@
|
||||
"""Parity tests: MLX refiner vs PyTorch reference (Phase 2).
|
||||
|
||||
Uses saved coarse predictions + RGB -> runs MLX refiner ->
|
||||
compares against saved PyTorch delta logits and final outputs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from corridorkey_mlx.model.refiner import CNNRefinerModule
|
||||
from corridorkey_mlx.utils.layout import conv_weight_pt_to_mlx, nchw_to_nhwc_np, nhwc_to_nchw_np
|
||||
|
||||
FIXTURE_PATH = Path("reference/fixtures/golden.npz")
|
||||
WEIGHTS_PATH = Path("reference/fixtures/golden_weights.npz")
|
||||
|
||||
MAX_ABS_TOL = 1e-4
|
||||
|
||||
|
||||
def _skip_if_missing() -> None:
|
||||
if not FIXTURE_PATH.exists() or not WEIGHTS_PATH.exists():
|
||||
pytest.skip("Fixture files not found — run dump_pytorch_reference.py first")
|
||||
|
||||
|
||||
def _load_refiner_weights(weights: dict[str, np.ndarray]) -> list[tuple[str, mx.array]]:
|
||||
"""Convert PyTorch refiner state_dict to MLX weight list.
|
||||
|
||||
Handles:
|
||||
- Conv weight transpose (O,I,H,W) -> (O,H,W,I)
|
||||
- Key mapping from PyTorch Sequential indices to named attributes
|
||||
"""
|
||||
prefix = "refiner."
|
||||
weight_list: list[tuple[str, mx.array]] = []
|
||||
|
||||
# Map PyTorch stem Sequential keys to our named attributes
|
||||
stem_key_map = {
|
||||
"stem.0.weight": "stem_conv.weight",
|
||||
"stem.0.bias": "stem_conv.bias",
|
||||
"stem.1.weight": "stem_gn.weight",
|
||||
"stem.1.bias": "stem_gn.bias",
|
||||
}
|
||||
|
||||
for pt_key, value in weights.items():
|
||||
if not pt_key.startswith(prefix):
|
||||
continue
|
||||
mlx_key = pt_key[len(prefix) :]
|
||||
|
||||
# Remap stem keys
|
||||
if mlx_key in stem_key_map:
|
||||
mlx_key = stem_key_map[mlx_key]
|
||||
|
||||
# Conv2d weights: (O,I,H,W) -> (O,H,W,I)
|
||||
if mlx_key.endswith(".weight") and _is_conv_weight(value):
|
||||
value = conv_weight_pt_to_mlx(value)
|
||||
|
||||
weight_list.append((mlx_key, mx.array(value)))
|
||||
|
||||
return weight_list
|
||||
|
||||
|
||||
def _is_conv_weight(value: np.ndarray) -> bool:
|
||||
"""Check if a weight tensor is a conv weight (4D with spatial dims)."""
|
||||
return value.ndim == 4
|
||||
|
||||
|
||||
def _build_and_load_refiner(weights: dict[str, np.ndarray]) -> CNNRefinerModule:
|
||||
"""Build an MLX CNNRefinerModule and load converted weights."""
|
||||
refiner = CNNRefinerModule()
|
||||
weight_list = _load_refiner_weights(weights)
|
||||
refiner.load_weights(weight_list)
|
||||
refiner.eval()
|
||||
mx.eval(refiner.parameters()) # noqa: S307 — mx.eval is MLX's compute trigger
|
||||
return refiner
|
||||
|
||||
|
||||
def test_refiner_delta_parity() -> None:
|
||||
"""MLX refiner delta logits match PyTorch within tolerance."""
|
||||
_skip_if_missing()
|
||||
|
||||
fixtures = dict(np.load(FIXTURE_PATH))
|
||||
weights = dict(np.load(WEIGHTS_PATH))
|
||||
|
||||
# RGB: first 3 channels of input
|
||||
rgb_nchw = fixtures["input"][:, :3] # (1, 3, H, W)
|
||||
rgb_nhwc = mx.array(nchw_to_nhwc_np(rgb_nchw))
|
||||
|
||||
# Coarse predictions: alpha_coarse (1ch) + fg_coarse (3ch) = 4ch
|
||||
alpha_coarse_nhwc = nchw_to_nhwc_np(fixtures["alpha_coarse"])
|
||||
fg_coarse_nhwc = nchw_to_nhwc_np(fixtures["fg_coarse"])
|
||||
coarse_pred = mx.array(np.concatenate([alpha_coarse_nhwc, fg_coarse_nhwc], axis=-1))
|
||||
|
||||
refiner = _build_and_load_refiner(weights)
|
||||
result_nhwc = refiner(rgb_nhwc, coarse_pred)
|
||||
mx.eval(result_nhwc) # noqa: S307
|
||||
|
||||
result_nchw = nhwc_to_nchw_np(np.array(result_nhwc))
|
||||
expected = fixtures["delta_logits"]
|
||||
|
||||
max_abs_err = float(np.max(np.abs(result_nchw - expected)))
|
||||
mean_abs_err = float(np.mean(np.abs(result_nchw - expected)))
|
||||
print(f"\nRefiner delta parity — max_abs: {max_abs_err:.6e}, mean_abs: {mean_abs_err:.6e}")
|
||||
|
||||
assert result_nchw.shape == expected.shape
|
||||
assert max_abs_err < MAX_ABS_TOL, (
|
||||
f"Max abs error {max_abs_err:.6e} exceeds tolerance {MAX_ABS_TOL}"
|
||||
)
|
||||
|
||||
|
||||
def test_refiner_final_output_parity() -> None:
|
||||
"""MLX final alpha and fg match PyTorch within tolerance.
|
||||
|
||||
Tests the full residual + sigmoid path using MLX refiner output.
|
||||
"""
|
||||
_skip_if_missing()
|
||||
|
||||
fixtures = dict(np.load(FIXTURE_PATH))
|
||||
weights = dict(np.load(WEIGHTS_PATH))
|
||||
|
||||
# Run refiner
|
||||
rgb_nchw = fixtures["input"][:, :3]
|
||||
rgb_nhwc = mx.array(nchw_to_nhwc_np(rgb_nchw))
|
||||
alpha_coarse_nhwc = nchw_to_nhwc_np(fixtures["alpha_coarse"])
|
||||
fg_coarse_nhwc = nchw_to_nhwc_np(fixtures["fg_coarse"])
|
||||
coarse_pred = mx.array(np.concatenate([alpha_coarse_nhwc, fg_coarse_nhwc], axis=-1))
|
||||
|
||||
refiner = _build_and_load_refiner(weights)
|
||||
delta_nhwc = refiner(rgb_nhwc, coarse_pred)
|
||||
mx.eval(delta_nhwc) # noqa: S307
|
||||
|
||||
# Apply residual in logit space + sigmoid (matching PyTorch forward pass)
|
||||
alpha_logits_up = mx.array(nchw_to_nhwc_np(fixtures["alpha_logits_up"]))
|
||||
fg_logits_up = mx.array(nchw_to_nhwc_np(fixtures["fg_logits_up"]))
|
||||
|
||||
alpha_final = mx.sigmoid(alpha_logits_up + delta_nhwc[..., 0:1])
|
||||
fg_final = mx.sigmoid(fg_logits_up + delta_nhwc[..., 1:4])
|
||||
mx.eval(alpha_final, fg_final) # noqa: S307
|
||||
|
||||
# Compare
|
||||
for name, result, expected_key in [
|
||||
("alpha_final", alpha_final, "alpha_final"),
|
||||
("fg_final", fg_final, "fg_final"),
|
||||
]:
|
||||
result_nchw = nhwc_to_nchw_np(np.array(result))
|
||||
expected = fixtures[expected_key]
|
||||
max_abs_err = float(np.max(np.abs(result_nchw - expected)))
|
||||
mean_abs_err = float(np.mean(np.abs(result_nchw - expected)))
|
||||
print(f"\n{name} parity — max_abs: {max_abs_err:.6e}, mean_abs: {mean_abs_err:.6e}")
|
||||
assert max_abs_err < MAX_ABS_TOL, (
|
||||
f"{name}: max abs error {max_abs_err:.6e} exceeds tolerance {MAX_ABS_TOL}"
|
||||
)
|
||||
@ -1,78 +0,0 @@
|
||||
"""2048 smoke tests — execution check at native resolution.
|
||||
|
||||
test_smoke_2048_wiring: lightweight, no checkpoint, runs in normal suite.
|
||||
test_smoke_2048_full: loads real checkpoint at 2048, marked slow + skipif.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from corridorkey_mlx.model.corridorkey import GreenFormer
|
||||
|
||||
CHECKPOINT = Path("checkpoints/corridorkey_mlx.safetensors")
|
||||
HAS_CHECKPOINT = CHECKPOINT.exists()
|
||||
|
||||
WIRING_IMG_SIZE = 256 # small enough to run fast with random weights
|
||||
|
||||
|
||||
def test_smoke_2048_wiring() -> None:
|
||||
"""GreenFormer constructs at 2048 and produces correct shapes (random weights)."""
|
||||
model = GreenFormer(img_size=WIRING_IMG_SIZE)
|
||||
x = mx.random.normal((1, WIRING_IMG_SIZE, WIRING_IMG_SIZE, 4))
|
||||
# mx.eval materializes lazy MLX arrays (not Python eval)
|
||||
mx.eval(x) # noqa: S307
|
||||
|
||||
out = model(x)
|
||||
mx.eval(out) # noqa: S307
|
||||
|
||||
assert out["alpha_final"].shape == (1, WIRING_IMG_SIZE, WIRING_IMG_SIZE, 1)
|
||||
assert out["fg_final"].shape == (1, WIRING_IMG_SIZE, WIRING_IMG_SIZE, 3)
|
||||
|
||||
# No NaN/Inf
|
||||
alpha = np.array(out["alpha_final"])
|
||||
fg = np.array(out["fg_final"])
|
||||
assert not np.isnan(alpha).any(), "alpha_final contains NaN"
|
||||
assert not np.isnan(fg).any(), "fg_final contains NaN"
|
||||
assert not np.isinf(alpha).any(), "alpha_final contains Inf"
|
||||
assert not np.isinf(fg).any(), "fg_final contains Inf"
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.skipif(not HAS_CHECKPOINT, reason="Checkpoint not available")
|
||||
def test_smoke_2048_full() -> None:
|
||||
"""Full 2048 inference with real checkpoint via engine."""
|
||||
from corridorkey_mlx import CorridorKeyMLXEngine
|
||||
|
||||
engine = CorridorKeyMLXEngine(
|
||||
checkpoint_path=CHECKPOINT,
|
||||
img_size=2048,
|
||||
compile=False, # skip compile overhead in test
|
||||
)
|
||||
|
||||
# Use synthetic inputs at 2048
|
||||
rng = np.random.default_rng(42)
|
||||
rgb = rng.integers(0, 256, (2048, 2048, 3), dtype=np.uint8)
|
||||
mask = rng.integers(0, 256, (2048, 2048), dtype=np.uint8)
|
||||
|
||||
result = engine.process_frame(rgb, mask)
|
||||
|
||||
# Shape checks
|
||||
assert result["alpha"].shape == (2048, 2048)
|
||||
assert result["alpha"].dtype == np.uint8
|
||||
assert result["fg"].shape == (2048, 2048, 3)
|
||||
assert result["fg"].dtype == np.uint8
|
||||
assert result["comp"].shape == (2048, 2048, 3)
|
||||
|
||||
# No NaN/Inf
|
||||
for key in ("alpha", "fg", "comp"):
|
||||
arr = result[key]
|
||||
assert not np.isnan(arr).any(), f"{key} contains NaN"
|
||||
assert not np.isinf(arr).any(), f"{key} contains Inf"
|
||||
|
||||
# Alpha shouldn't be degenerate
|
||||
assert result["alpha"].min() != result["alpha"].max(), "alpha is constant — suspicious"
|
||||
@ -1,15 +1,9 @@
|
||||
"""Unit tests for weights download module.
|
||||
|
||||
Tests config, caching, checksum verification, and CLI arg parsing
|
||||
without requiring real network access.
|
||||
"""
|
||||
"""Tests for weights download module: config defaults, checksum, hash parsing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
from typing import TYPE_CHECKING
|
||||
from unittest.mock import patch
|
||||
|
||||
from corridorkey_mlx.weights import (
|
||||
DEFAULT_ASSET_NAME,
|
||||
@ -22,7 +16,6 @@ from corridorkey_mlx.weights import (
|
||||
github_owner_repo,
|
||||
verify_sha256,
|
||||
)
|
||||
from corridorkey_mlx.weights_cli import build_parser
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
@ -37,26 +30,12 @@ class TestConfig:
|
||||
assert owner == DEFAULT_GITHUB_OWNER
|
||||
assert repo == DEFAULT_GITHUB_REPO
|
||||
|
||||
def test_owner_repo_env_override(self) -> None:
|
||||
with patch.dict(os.environ, {"CORRIDORKEY_MLX_WEIGHTS_REPO": "other/repo"}):
|
||||
owner, repo = github_owner_repo()
|
||||
assert owner == "other"
|
||||
assert repo == "repo"
|
||||
|
||||
def test_default_asset_name(self) -> None:
|
||||
assert default_asset_name() == DEFAULT_ASSET_NAME
|
||||
|
||||
def test_asset_name_env_override(self) -> None:
|
||||
with patch.dict(os.environ, {"CORRIDORKEY_MLX_WEIGHTS_ASSET": "custom.st"}):
|
||||
assert default_asset_name() == "custom.st"
|
||||
|
||||
def test_default_tag(self) -> None:
|
||||
assert default_tag() == "latest"
|
||||
|
||||
def test_tag_env_override(self) -> None:
|
||||
with patch.dict(os.environ, {"CORRIDORKEY_MLX_WEIGHTS_TAG": "v2.0"}):
|
||||
assert default_tag() == "v2.0"
|
||||
|
||||
|
||||
# -- Cache directory ---------------------------------------------------------
|
||||
|
||||
@ -91,45 +70,14 @@ class TestChecksum:
|
||||
def test_parse_hash_garbage(self) -> None:
|
||||
assert _parse_hash_line("not a hash", "file.bin") is None
|
||||
|
||||
def test_verify_sha256(self, tmp_path: Path) -> None:
|
||||
def test_verify_sha256_valid(self, tmp_path: Path) -> None:
|
||||
content = b"hello world"
|
||||
expected = hashlib.sha256(content).hexdigest()
|
||||
p = tmp_path / "test.bin"
|
||||
p.write_bytes(content)
|
||||
|
||||
assert verify_sha256(p, expected) is True
|
||||
|
||||
def test_verify_sha256_invalid(self, tmp_path: Path) -> None:
|
||||
p = tmp_path / "test.bin"
|
||||
p.write_bytes(b"hello world")
|
||||
assert verify_sha256(p, "0" * 64) is False
|
||||
|
||||
|
||||
# -- CLI arg parsing ---------------------------------------------------------
|
||||
|
||||
|
||||
class TestCLI:
|
||||
def test_download_defaults(self) -> None:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args(["download"])
|
||||
assert args.command == "download"
|
||||
assert args.tag is None
|
||||
assert args.asset_name is None
|
||||
assert args.out is None
|
||||
assert args.force is False
|
||||
assert args.no_verify is False
|
||||
assert args.print_path is False
|
||||
|
||||
def test_download_all_flags(self) -> None:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args([
|
||||
"download",
|
||||
"--tag", "v1.0",
|
||||
"--asset", "custom.st",
|
||||
"--out", "/tmp/w",
|
||||
"--force",
|
||||
"--no-verify",
|
||||
"--print-path",
|
||||
])
|
||||
assert args.tag == "v1.0"
|
||||
assert args.asset_name == "custom.st"
|
||||
assert args.out == "/tmp/w"
|
||||
assert args.force is True
|
||||
assert args.no_verify is True
|
||||
assert args.print_path is True
|
||||
|
||||
Loading…
Reference in New Issue
Block a user