From 9b079616190cd13838bff07c64197506f301132a Mon Sep 17 00:00:00 2001 From: cmoyates Date: Sun, 1 Mar 2026 04:49:01 -0330 Subject: [PATCH] feat(phase1): PyTorch reference harness + fixture tests 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 --- scripts/dump_pytorch_reference.py | 425 ++++++++++++++++++++++++++++++ tests/test_reference_fixtures.py | 130 +++++++++ 2 files changed, 555 insertions(+) create mode 100644 scripts/dump_pytorch_reference.py create mode 100644 tests/test_reference_fixtures.py diff --git a/scripts/dump_pytorch_reference.py b/scripts/dump_pytorch_reference.py new file mode 100644 index 0000000..41f803e --- /dev/null +++ b/scripts/dump_pytorch_reference.py @@ -0,0 +1,425 @@ +#!/usr/bin/env python3 +"""Dump intermediate PyTorch tensors for MLX parity testing. + +Loads original CorridorKey checkpoint, runs a forward pass, +and saves intermediate activations to reference/fixtures/. + +Usage: + uv run --group reference python scripts/dump_pytorch_reference.py \ + --checkpoint checkpoints/CorridorKey_v1.0.pth + +Requires: torch, timm, rich (install via `uv sync --group reference`) +""" + +from __future__ import annotations + +import argparse +from collections import OrderedDict +from pathlib import Path + +import numpy as np +import timm +import torch +import torch.nn as nn +import torch.nn.functional as F +from rich.console import Console +from rich.table import Table + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- +SEED = 42 +IMG_SIZE = 512 +BACKBONE_CHANNELS = [112, 224, 448, 896] +EMBED_DIM = 256 +REFINER_CHANNELS = 64 +REFINER_GROUPS = 8 +REFINER_SCALE = 10.0 +DROPOUT_RATE = 0.1 +HIERA_MODEL_NAME = "hiera_base_plus_224.mae_in1k_ft_in1k" +INPUT_CHANNELS = 4 # RGB + alpha hint +OUTPUT_DIR = Path("reference/fixtures") +FIXTURE_FILENAME = "golden.npz" + +console = Console() + + +# --------------------------------------------------------------------------- +# Model components (mirror of nikopueringer/CorridorKey GreenFormer) +# --------------------------------------------------------------------------- +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 forward(self, x: torch.Tensor) -> torch.Tensor: + return self.proj(x) + + +class DecoderHead(nn.Module): + """SegFormer-style multiscale feature fusion head.""" + + 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.BatchNorm2d(embed_dim) + self.classifier = nn.Conv2d(embed_dim, output_dim, kernel_size=1) + + def forward(self, features: list[torch.Tensor]) -> torch.Tensor: + c1, c2, c3, c4 = features + target_size = 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, c, h, w = feat.shape + # Flatten spatial -> project -> reshape + x = feat.flatten(2).transpose(1, 2) # [B, H*W, C] + x = linear(x) # [B, H*W, embed_dim] + x = x.transpose(1, 2).reshape(b, -1, h, w) # [B, embed_dim, H, W] + x = F.interpolate(x, size=target_size, mode="bilinear", align_corners=False) + projected.append(x) + + fused = torch.cat(projected, dim=1) # [B, embed_dim*4, H/4, W/4] + fused = self.linear_fuse(fused) + fused = self.bn(fused) + fused = F.relu(fused) + fused = F.dropout(fused, p=DROPOUT_RATE, training=self.training) + return self.classifier(fused) + + +class RefinerBlock(nn.Module): + """Dilated residual block with GroupNorm.""" + + 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) + self.conv2 = nn.Conv2d( + channels, channels, kernel_size=3, padding=dilation, dilation=dilation + ) + self.gn2 = nn.GroupNorm(REFINER_GROUPS, channels) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + out = F.relu(self.gn1(self.conv1(x))) + out = self.gn2(self.conv2(out)) + return F.relu(out + residual) + + +class CNNRefinerModule(nn.Module): + """CNN refiner: stem + 4 dilated ResBlocks + 1x1 projection.""" + + def __init__(self) -> None: + super().__init__() + # 7 input channels: RGB (3) + coarse_pred (4: alpha + fg) + refiner_input_channels = 7 + self.stem = nn.Sequential( + nn.Conv2d(refiner_input_channels, REFINER_CHANNELS, kernel_size=3, padding=1), + nn.GroupNorm(REFINER_GROUPS, REFINER_CHANNELS), + nn.ReLU(), + ) + 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) + # 4 output channels: delta for alpha (1) + delta for fg (3) + refiner_output_channels = 4 + self.final = nn.Conv2d(REFINER_CHANNELS, refiner_output_channels, kernel_size=1) + + def forward(self, rgb: torch.Tensor, coarse_pred: torch.Tensor) -> torch.Tensor: + x = torch.cat([rgb, coarse_pred], dim=1) # [B, 7, H, W] + x = self.stem(x) + x = self.res1(x) + x = self.res2(x) + x = self.res3(x) + x = self.res4(x) + return self.final(x) * REFINER_SCALE + + +class GreenFormer(nn.Module): + """Top-level CorridorKey model: Hiera encoder + dual decoder heads + CNN refiner.""" + + def __init__(self) -> None: + super().__init__() + self.encoder = timm.create_model( + HIERA_MODEL_NAME, + pretrained=False, + features_only=True, + img_size=IMG_SIZE, + ) + self._patch_first_conv() + + self.alpha_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=1) + self.fg_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=3) + self.refiner = CNNRefinerModule() + + def _patch_first_conv(self) -> None: + """Replace 3-channel patch embed conv with 4-channel version.""" + old_conv = self.encoder.model.patch_embed.proj + new_conv = nn.Conv2d( + INPUT_CHANNELS, + old_conv.out_channels, + kernel_size=old_conv.kernel_size, + stride=old_conv.stride, + padding=old_conv.padding, + bias=old_conv.bias is not None, + ) + # Zero-init; checkpoint weights will overwrite + nn.init.zeros_(new_conv.weight) + if old_conv.bias is not None: + new_conv.bias = old_conv.bias + new_conv.weight.data[:, :3] = old_conv.weight.data + self.encoder.model.patch_embed.proj = new_conv + + def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]: + input_size = x.shape[2:] # (H, W) + + # Backbone + features = self.encoder(x) # list of 4 feature maps + + # Decoder heads (at H/4 resolution) + alpha_logits = self.alpha_decoder(features) + fg_logits = self.fg_decoder(features) + + # Upsample to input resolution + alpha_logits_up = F.interpolate( + alpha_logits, size=input_size, mode="bilinear", align_corners=False + ) + fg_logits_up = F.interpolate( + fg_logits, size=input_size, mode="bilinear", align_corners=False + ) + + # Coarse predictions + alpha_coarse = torch.sigmoid(alpha_logits_up) + fg_coarse = torch.sigmoid(fg_logits_up) + + # Refiner + rgb = x[:, :3] + coarse_pred = torch.cat([alpha_coarse, fg_coarse], dim=1) # [B, 4, H, W] + delta_logits = self.refiner(rgb, coarse_pred) + + # Final predictions (residual in logit space) + alpha_final = torch.sigmoid(alpha_logits_up + delta_logits[:, 0:1]) + fg_final = torch.sigmoid(fg_logits_up + delta_logits[:, 1:4]) + + return { + "encoder_features": features, + "alpha_logits": alpha_logits, + "fg_logits": fg_logits, + "alpha_logits_up": alpha_logits_up, + "fg_logits_up": fg_logits_up, + "alpha_coarse": alpha_coarse, + "fg_coarse": fg_coarse, + "delta_logits": delta_logits, + "alpha_final": alpha_final, + "fg_final": fg_final, + } + + +# --------------------------------------------------------------------------- +# Checkpoint loading +# --------------------------------------------------------------------------- +def _interpolate_pos_embed( + ckpt_embed: torch.Tensor, + model_embed: torch.Tensor, + embed_dim: int, +) -> torch.Tensor: + """Bicubic interpolation of pos_embed from checkpoint resolution to model resolution. + + Both tensors are shape (1, N, C) where N = H*W spatial tokens. + """ + ckpt_n = ckpt_embed.shape[1] + model_n = model_embed.shape[1] + if ckpt_n == model_n: + return ckpt_embed + + import math + + ckpt_side = int(math.sqrt(ckpt_n)) + model_side = int(math.sqrt(model_n)) + assert ckpt_side * ckpt_side == ckpt_n, f"pos_embed not square: {ckpt_n}" + assert model_side * model_side == model_n, f"pos_embed not square: {model_n}" + + console.print( + f"[cyan]Interpolating pos_embed: {ckpt_side}x{ckpt_side} -> " + f"{model_side}x{model_side}[/cyan]" + ) + + # (1, N, C) -> (1, C, H, W) for interpolation + embed = ckpt_embed.reshape(1, ckpt_side, ckpt_side, embed_dim).permute(0, 3, 1, 2) + embed = F.interpolate( + embed, size=(model_side, model_side), mode="bicubic", align_corners=False + ) + # Back to (1, N, C) + return embed.permute(0, 2, 3, 1).reshape(1, model_n, embed_dim) + + +def load_checkpoint(model: GreenFormer, checkpoint_path: Path) -> None: + """Load state_dict from checkpoint, stripping torch.compile prefix if present.""" + raw = torch.load(checkpoint_path, map_location="cpu", weights_only=True) + state_dict = raw.get("state_dict", raw) + + # Strip _orig_mod. prefix from torch.compile + compile_prefix = "_orig_mod." + cleaned: OrderedDict[str, torch.Tensor] = OrderedDict() + for key, value in state_dict.items(): + clean_key = key[len(compile_prefix) :] if key.startswith(compile_prefix) else key + cleaned[clean_key] = value + + # Interpolate pos_embed if spatial dimensions differ + pos_embed_key = "encoder.model.pos_embed" + if pos_embed_key in cleaned: + model_embed = model.state_dict()[pos_embed_key] + if cleaned[pos_embed_key].shape != model_embed.shape: + cleaned[pos_embed_key] = _interpolate_pos_embed( + cleaned[pos_embed_key], + model_embed, + embed_dim=BACKBONE_CHANNELS[0], + ) + + missing, unexpected = model.load_state_dict(cleaned, strict=False) + if missing: + console.print(f"[yellow]Missing keys ({len(missing)}):[/yellow]") + for k in missing: + console.print(f" {k}") + if unexpected: + console.print(f"[yellow]Unexpected keys ({len(unexpected)}):[/yellow]") + for k in unexpected: + console.print(f" {k}") + if not missing and not unexpected: + console.print("[green]All keys matched.[/green]") + + +# --------------------------------------------------------------------------- +# Tensor dumping +# --------------------------------------------------------------------------- +def dump_fixtures(outputs: dict[str, torch.Tensor | list[torch.Tensor]], output_dir: Path) -> None: + """Save all intermediate tensors to a single .npz file.""" + output_dir.mkdir(parents=True, exist_ok=True) + + arrays: dict[str, np.ndarray] = {} + + # Flatten encoder features into separate keys + features = outputs.pop("encoder_features") + for i, feat in enumerate(features): + arrays[f"encoder_feature_{i}"] = feat.cpu().numpy() + + # Remaining tensors + for name, tensor in outputs.items(): + assert isinstance(tensor, torch.Tensor), f"Expected Tensor for {name}, got {type(tensor)}" + arrays[name] = tensor.cpu().numpy() + + out_path = output_dir / FIXTURE_FILENAME + np.savez(out_path, **arrays) + console.print(f"\n[green]Saved fixtures 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") + table.add_column("Tensor", style="cyan") + table.add_column("Shape", style="green") + table.add_column("Dtype", style="yellow") + table.add_column("Min", style="dim") + table.add_column("Max", style="dim") + + for name, value in outputs.items(): + if isinstance(value, list): + for i, v in enumerate(value): + table.add_row( + f"{name}[{i}]", + str(tuple(v.shape)), + str(v.dtype), + f"{v.min().item():.4f}", + f"{v.max().item():.4f}", + ) + else: + table.add_row( + name, + str(tuple(value.shape)), + str(value.dtype), + f"{value.min().item():.4f}", + f"{value.max().item():.4f}", + ) + + console.print(table) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- +def main() -> None: + parser = argparse.ArgumentParser(description="Dump PyTorch reference fixtures") + parser.add_argument( + "--checkpoint", + type=Path, + default=Path("checkpoints/CorridorKey_v1.0.pth"), + help="Path to CorridorKey checkpoint", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=OUTPUT_DIR, + help="Directory for fixture output", + ) + parser.add_argument( + "--seed", + type=int, + default=SEED, + help="Random seed for deterministic input", + ) + args = parser.parse_args() + + if not args.checkpoint.exists(): + console.print(f"[red]Checkpoint not found: {args.checkpoint}[/red]") + console.print("Download CorridorKey_v1.0.pth and place it in checkpoints/") + raise SystemExit(1) + + # Deterministic setup + torch.manual_seed(args.seed) + torch.set_grad_enabled(False) + + console.print("[bold]Building GreenFormer model...[/bold]") + model = GreenFormer() + # Set to evaluation mode (disables dropout, uses running stats for BatchNorm) + model.train(False) + + console.print(f"[bold]Loading checkpoint: {args.checkpoint}[/bold]") + load_checkpoint(model, args.checkpoint) + + # Deterministic input + torch.manual_seed(args.seed) + sample_input = torch.randn(1, INPUT_CHANNELS, IMG_SIZE, IMG_SIZE) + + shape_str = str(tuple(sample_input.shape)) + console.print(f"[bold]Running forward pass (input shape: {shape_str})...[/bold]") + outputs = model(sample_input) + + # Add input to outputs for completeness + outputs = {"input": sample_input, **outputs} + + print_shape_report(outputs) + dump_fixtures(outputs, args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/tests/test_reference_fixtures.py b/tests/test_reference_fixtures.py new file mode 100644 index 0000000..5ce3a42 --- /dev/null +++ b/tests/test_reference_fixtures.py @@ -0,0 +1,130 @@ +"""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