9b07961619
Dump script loads CorridorKey checkpoint, runs deterministic forward pass (seed=42, 512x512), saves 14 intermediate tensors to golden.npz. Includes pos_embed bicubic interpolation for resolution mismatch. 32 shape/dtype/range validation tests all passing. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
131 lines
4.6 KiB
Python
131 lines
4.6 KiB
Python
"""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
|