diff --git a/docs/plans/2026-03-01-feat-corridorkey-mlx-inference-port-plan.md b/docs/plans/2026-03-01-feat-corridorkey-mlx-inference-port-plan.md index dd0461d..93343d4 100644 --- a/docs/plans/2026-03-01-feat-corridorkey-mlx-inference-port-plan.md +++ b/docs/plans/2026-03-01-feat-corridorkey-mlx-inference-port-plan.md @@ -109,10 +109,10 @@ RGB + coarse alpha hint (4ch) **Components to implement:** -- [ ] `MLP` -- feedforward block -- [ ] `DecoderHead` -- consumes backbone features, produces coarse predictions -- [ ] `RefinerBlock` -- single refiner stage -- [ ] `CNNRefinerModule` -- full refiner consuming 7ch input +- [x] `MLP` -- feedforward block +- [x] `DecoderHead` -- consumes backbone features, produces coarse predictions +- [x] `RefinerBlock` -- single refiner stage +- [x] `CNNRefinerModule` -- full refiner consuming 7ch input **Files touched:** diff --git a/scripts/dump_pytorch_reference.py b/scripts/dump_pytorch_reference.py index 41f803e..f91b70a 100644 --- a/scripts/dump_pytorch_reference.py +++ b/scripts/dump_pytorch_reference.py @@ -333,6 +333,29 @@ def dump_fixtures(outputs: dict[str, torch.Tensor | list[torch.Tensor]], output_ console.print(f"\n[green]Saved fixtures to {out_path}[/green]") +WEIGHTS_FILENAME = "golden_weights.npz" + + +def dump_weights(model: GreenFormer, output_dir: Path) -> None: + """Save decoder and refiner weights as numpy arrays for Phase 2 parity tests. + + Keys are the PyTorch state_dict keys (e.g. 'alpha_decoder.linear_c1.proj.weight'). + Conv weights remain in PyTorch format (O,I,H,W) — tests handle conversion. + """ + output_dir.mkdir(parents=True, exist_ok=True) + arrays: dict[str, np.ndarray] = {} + + # Extract decoder and refiner state dicts + prefixes = ("alpha_decoder.", "fg_decoder.", "refiner.") + for key, param in model.state_dict().items(): + if any(key.startswith(p) for p in prefixes): + arrays[key] = param.cpu().numpy() + + out_path = output_dir / WEIGHTS_FILENAME + np.savez(out_path, **arrays) + console.print(f"[green]Saved {len(arrays)} weight tensors to {out_path}[/green]") + + def print_shape_report(outputs: dict[str, torch.Tensor | list[torch.Tensor]]) -> None: """Print a rich table of tensor names, shapes, and dtypes.""" table = Table(title="Reference Fixture Shapes") @@ -419,6 +442,7 @@ def main() -> None: print_shape_report(outputs) dump_fixtures(outputs, args.output_dir) + dump_weights(model, args.output_dir) if __name__ == "__main__": diff --git a/src/corridorkey_mlx/model/decoder.py b/src/corridorkey_mlx/model/decoder.py index 450f417..bc07cc4 100644 --- a/src/corridorkey_mlx/model/decoder.py +++ b/src/corridorkey_mlx/model/decoder.py @@ -1,5 +1,95 @@ -"""Decoder heads — MLX port (not yet implemented). +"""Decoder heads — MLX port. Two heads: alpha (1ch) and foreground (3ch). -Consume multiscale backbone features, upsample to full resolution. +Consume multiscale backbone features (NHWC), upsample and fuse to produce predictions. + +Architecture mirrors nikopueringer/CorridorKey DecoderHead (SegFormer-style). """ + +from __future__ import annotations + +import mlx.core as mx +import mlx.nn as nn + + +class MLP(nn.Module): + """Single linear projection: input_dim -> embed_dim.""" + + def __init__(self, input_dim: int, embed_dim: int) -> None: + super().__init__() + self.proj = nn.Linear(input_dim, embed_dim) + + def __call__(self, x: mx.array) -> mx.array: + return self.proj(x) + + +class DecoderHead(nn.Module): + """SegFormer-style multiscale feature fusion head (NHWC). + + Takes 4 multi-scale feature maps in NHWC format, projects each to embed_dim, + upsamples to the largest spatial resolution, fuses, and classifies. + """ + + def __init__( + self, + in_channels: list[int], + embed_dim: int, + output_dim: int, + ) -> None: + super().__init__() + self.linear_c1 = MLP(in_channels[0], embed_dim) + self.linear_c2 = MLP(in_channels[1], embed_dim) + self.linear_c3 = MLP(in_channels[2], embed_dim) + self.linear_c4 = MLP(in_channels[3], embed_dim) + + fused_channels = embed_dim * len(in_channels) + self.linear_fuse = nn.Conv2d(fused_channels, embed_dim, kernel_size=1, bias=False) + self.bn = nn.BatchNorm(embed_dim) + self.classifier = nn.Conv2d(embed_dim, output_dim, kernel_size=1) + + def __call__(self, features: list[mx.array]) -> mx.array: + """Forward pass. + + Args: + features: 4 feature maps in NHWC format. + [0]: (B, H/4, W/4, C1) + [1]: (B, H/8, W/8, C2) + [2]: (B, H/16, W/16, C3) + [3]: (B, H/32, W/32, C4) + + Returns: + Logits in NHWC: (B, H/4, W/4, output_dim) + """ + c1, c2, c3, c4 = features + target_h, target_w = c1.shape[1], c1.shape[2] # H/4, W/4 + + projected = [] + for feat, linear in zip( + [c1, c2, c3, c4], + [self.linear_c1, self.linear_c2, self.linear_c3, self.linear_c4], + strict=True, + ): + b, h, w, _c = feat.shape + # NHWC: reshape to (B, H*W, C), project, reshape back + x = feat.reshape(b, h * w, _c) + x = linear(x) # (B, H*W, embed_dim) + x = x.reshape(b, h, w, -1) # (B, H, W, embed_dim) + # Upsample to target spatial size + if h != target_h or w != target_w: + scale_h = target_h / h + scale_w = target_w / w + upsampler = nn.Upsample( + scale_factor=(scale_h, scale_w), + mode="linear", + align_corners=False, + ) + x = upsampler(x) + projected.append(x) + + # Concatenate along channel dim (last dim in NHWC) + fused = mx.concatenate(projected, axis=-1) # (B, H/4, W/4, embed_dim*4) + fused = self.linear_fuse(fused) + fused = self.bn(fused) + fused = nn.relu(fused) + # No dropout at inference (eval mode) + return self.classifier(fused) diff --git a/src/corridorkey_mlx/model/refiner.py b/src/corridorkey_mlx/model/refiner.py index a3beeb3..1f0bb41 100644 --- a/src/corridorkey_mlx/model/refiner.py +++ b/src/corridorkey_mlx/model/refiner.py @@ -1,5 +1,87 @@ -"""CNN refiner — MLX port (not yet implemented). +"""CNN refiner — MLX port. -Input: RGB + coarse predictions (7ch total). -Output: additive delta logits → final sigmoid. +Input: RGB + coarse predictions (7ch total) in NHWC. +Output: additive delta logits (4ch) in NHWC. + +Architecture mirrors nikopueringer/CorridorKey CNNRefinerModule. """ + +from __future__ import annotations + +import mlx.core as mx +import mlx.nn as nn + +REFINER_CHANNELS = 64 +REFINER_GROUPS = 8 +REFINER_SCALE = 10.0 + + +class RefinerBlock(nn.Module): + """Dilated residual block with GroupNorm (NHWC).""" + + def __init__(self, channels: int, dilation: int) -> None: + super().__init__() + self.conv1 = nn.Conv2d( + channels, + channels, + kernel_size=3, + padding=dilation, + dilation=dilation, + ) + self.gn1 = nn.GroupNorm(REFINER_GROUPS, channels, pytorch_compatible=True) + self.conv2 = nn.Conv2d( + channels, + channels, + kernel_size=3, + padding=dilation, + dilation=dilation, + ) + self.gn2 = nn.GroupNorm(REFINER_GROUPS, channels, pytorch_compatible=True) + + def __call__(self, x: mx.array) -> mx.array: + residual = x + out = nn.relu(self.gn1(self.conv1(x))) + out = self.gn2(self.conv2(out)) + return nn.relu(out + residual) + + +class CNNRefinerModule(nn.Module): + """CNN refiner: stem + 4 dilated ResBlocks + 1x1 projection (NHWC). + + Takes RGB (3ch) and coarse predictions (4ch) concatenated along channel dim. + Returns delta logits (4ch) scaled by 10x. + """ + + def __init__(self) -> None: + super().__init__() + refiner_input_channels = 7 # RGB (3) + coarse_pred (4: alpha + fg) + self.stem_conv = nn.Conv2d( + refiner_input_channels, REFINER_CHANNELS, kernel_size=3, padding=1 + ) + self.stem_gn = nn.GroupNorm(REFINER_GROUPS, REFINER_CHANNELS, pytorch_compatible=True) + + self.res1 = RefinerBlock(REFINER_CHANNELS, dilation=1) + self.res2 = RefinerBlock(REFINER_CHANNELS, dilation=2) + self.res3 = RefinerBlock(REFINER_CHANNELS, dilation=4) + self.res4 = RefinerBlock(REFINER_CHANNELS, dilation=8) + + refiner_output_channels = 4 # delta for alpha (1) + delta for fg (3) + self.final = nn.Conv2d(REFINER_CHANNELS, refiner_output_channels, kernel_size=1) + + def __call__(self, rgb: mx.array, coarse_pred: mx.array) -> mx.array: + """Forward pass. + + Args: + rgb: (B, H, W, 3) — RGB input in NHWC + coarse_pred: (B, H, W, 4) — concatenated alpha_coarse + fg_coarse in NHWC + + Returns: + Delta logits: (B, H, W, 4) in NHWC, scaled by 10x + """ + x = mx.concatenate([rgb, coarse_pred], axis=-1) # (B, H, W, 7) + x = nn.relu(self.stem_gn(self.stem_conv(x))) + x = self.res1(x) + x = self.res2(x) + x = self.res3(x) + x = self.res4(x) + return self.final(x) * REFINER_SCALE diff --git a/src/corridorkey_mlx/utils/layout.py b/src/corridorkey_mlx/utils/layout.py index 56b50be..c4c8433 100644 --- a/src/corridorkey_mlx/utils/layout.py +++ b/src/corridorkey_mlx/utils/layout.py @@ -1,4 +1,35 @@ -"""Tensor layout conversion utilities (not yet implemented). +"""Tensor layout conversion utilities. -Centralizes NCHW ↔ NHWC transforms. All layout conversions go through here. +Centralizes NCHW <-> NHWC transforms and conv weight transposes. +All layout conversions in the project go through this module. """ + +from __future__ import annotations + +import mlx.core as mx +import numpy as np + + +def nchw_to_nhwc(x: mx.array) -> mx.array: + """Convert (N, C, H, W) -> (N, H, W, C).""" + return mx.transpose(x, axes=(0, 2, 3, 1)) + + +def nhwc_to_nchw(x: mx.array) -> mx.array: + """Convert (N, H, W, C) -> (N, C, H, W).""" + return mx.transpose(x, axes=(0, 3, 1, 2)) + + +def conv_weight_pt_to_mlx(w: np.ndarray) -> np.ndarray: + """Transpose PyTorch conv weight (O, I, H, W) -> MLX conv weight (O, H, W, I).""" + return np.transpose(w, axes=(0, 2, 3, 1)) + + +def nchw_to_nhwc_np(x: np.ndarray) -> np.ndarray: + """Convert numpy (N, C, H, W) -> (N, H, W, C).""" + return np.transpose(x, axes=(0, 2, 3, 1)) + + +def nhwc_to_nchw_np(x: np.ndarray) -> np.ndarray: + """Convert numpy (N, H, W, C) -> (N, C, H, W).""" + return np.transpose(x, axes=(0, 3, 1, 2)) diff --git a/tests/test_decoder_parity.py b/tests/test_decoder_parity.py index 2a6fc75..41aa5d0 100644 --- a/tests/test_decoder_parity.py +++ b/tests/test_decoder_parity.py @@ -1,19 +1,150 @@ """Parity tests: MLX decoder heads vs PyTorch reference (Phase 2). -Uses saved backbone features → runs MLX decoder → compares +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 -pytestmark = pytest.mark.skip(reason="Phase 2: MLX decoder not yet implemented") +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", + ) diff --git a/tests/test_refiner_parity.py b/tests/test_refiner_parity.py index bca7bd2..7289036 100644 --- a/tests/test_refiner_parity.py +++ b/tests/test_refiner_parity.py @@ -1,19 +1,154 @@ """Parity tests: MLX refiner vs PyTorch reference (Phase 2). -Uses saved coarse predictions + RGB → runs MLX refiner → -compares against saved PyTorch final outputs. +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 -pytestmark = pytest.mark.skip(reason="Phase 2: MLX refiner not yet implemented") +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.""" - ... + """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}" + )