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 <noreply@anthropic.com>
This commit is contained in:
cmoyates
parent
897242d200
commit
9b07961619
425
scripts/dump_pytorch_reference.py
Normal file
425
scripts/dump_pytorch_reference.py
Normal file
@ -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()
|
||||
130
tests/test_reference_fixtures.py
Normal file
130
tests/test_reference_fixtures.py
Normal file
@ -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
|
||||
Loading…
Reference in New Issue
Block a user