diff --git a/README.md b/README.md index 7b58605..4493900 100644 --- a/README.md +++ b/README.md @@ -69,10 +69,11 @@ Inputs are 3D NIfTI volumes, all on the same voxel grid. Each run writes ## Technical details Preprocessing, sliding-window inference and postprocessing follow the pinned MONAI -1.6.0 bundles. Networks and runners are checked against seeded upstream recordings; -saved masks match exactly. Real-checkpoint parity and clinical performance are not -yet qualified. The lower-level API, tolerances and fixture regeneration are in -[docs/technical.md](docs/technical.md). +1.6.0 bundles. Networks and runners are checked against seeded upstream recordings. +The released renal checkpoint also matches the independent MONAI graph exactly +on the retained three-phase CT case, including raw logits and the saved mask. +Clinical performance remains unqualified. The lower-level API, tolerances and +fixture regeneration are in [docs/technical.md](docs/technical.md). ## Intended use diff --git a/docs/technical.md b/docs/technical.md index 9c75cc8..3b68d1d 100644 --- a/docs/technical.md +++ b/docs/technical.md @@ -19,17 +19,24 @@ MRI. Runtime code uses `medmlx-core`, NumPy and nibabel. Torch and MONAI are use only for conversion or reference generation. Importing the package or its model declaration does not import MLX. -**Qualification:** The recorded reference is the original MONAI implementation, -with synthetic seeded inputs and weights. Renal normalized intensities are -checked within one float32 rounding unit on [0, 1]; BraTS normalized intensities -within `2**-20` (about 9.54e-7). All saved masks agree exactly. The reviewer's -Metal run passed the seeded public-network, operator, sliding-window and runner -logit checks at a 1e-6 absolute bound. Its sole failure was an encoder -intermediate: 1.66893e-6 (seven float32 ULPs). Encoder intermediates now use a -`2**-19` absolute bound (about 1.91e-6), reflecting native accumulation rounding; -that changed assertion passes the current Metal suite (99 tests passed, two -released-checkpoint tests skipped). Real-checkpoint numerical parity, -clinical accuracy and full production-topology numerical parity remain unqualified. +**Qualification:** The reference is the original MONAI implementation. Seeded +public-network, operator, sliding-window and runner logit checks retain their +1e-6 absolute bound; encoder intermediates retain `2**-19`. Renal normalized +intensities are checked within one float32 rounding unit on [0, 1], and BraTS +normalized intensities within `2**-20` (about 9.54e-7). Saved fixture masks agree +exactly. + +The released renal checkpoint is also checked on dense and sparse seeded inputs +with unchanged `rtol=atol=1e-4` and exact sign agreement. On 12 October 2026, +source-order FP32 convolution and GroupNorm removed the previously retained +sparse/full-case discrepancy. On the retained external three-phase CT case, +all 85,089,840 raw logits were bitwise identical to a fresh independently parsed +MONAI CPU bundle run, and all 5,935,600 exported mask voxels agreed. This was +measured on Apple M1 Max, macOS 27.0.1, MLX 0.32.3, Torch 2.14.1 and MONAI 1.6.0; +bitwise equality is an observation for those runs, not a promise for every +platform or input. The case has no reference labels and retains manual +registration/cropping and a denoising substitution. Anatomical and clinical +accuracy remain unqualified. ## Install @@ -60,7 +67,11 @@ inference sections committed in `tests/fixtures/bundle_configs/` match those original bundle files. The network uses their inference graph in float32 with dropout disabled, group normalization, ReLU, skip addition, and either transposed convolution or trilinear 2x interpolation (`align_corners=False`). Native MLX -operations come from the shared core; no custom decoder kernels are required. +operations come from the shared core. Convolution and GroupNorm use its +source-order FP32 operators: channel-major fused multiply-add accumulation and +bounded four-lane Welford moments. Correcting only one of these operators leaves +amplified drift in the released renal checkpoint. The decoder uses the core's +phase-packed transposed convolution or trilinear interpolation. The BraTS bundle enables mixed-precision evaluation upstream. The reference contract uses its float32 evaluation graph; mixed-precision arithmetic is @@ -160,7 +171,9 @@ qform/sform agreement rules or rewrite input files. Renal preprocessing orders artery/vein/excretory channels, reorients to LPS, resamples to 0.8 mm, and clips/scales [-1000, 1000] HU to [0, 1]. Postprocessing inverts the spatial transforms, applies sigmoid and threshold `>= 0.5`, then -merges classes with the lowest nonzero channel winning. Labels 1 artery, 2 vein, +merges classes with the lowest nonzero channel winning. Sigmoid evaluates its +negative tail without exponent overflow while preserving the source float32 +expression near the 0.5 threshold. Labels 1 artery, 2 vein, 3 ureter, 4 neoplasm and 5 parenchyma are saved as float32. BraTS normalizes nonzero voxels per modality, applies sigmoid, inversion and @@ -224,4 +237,16 @@ writer, provenance, conversion-property and telemetry tests can run without it. Use `-m 'not mlx'` with pytest to run that portion of the suite in a sandbox without Metal. Full-suite collection does not load MLX; selecting a Metal test loads the shared runtime when the test executes. -There is no configured pyright check in this repository. +The strict Pyright and annotation checks are configured in `make typecheck`. +To include the released-checkpoint regression tests, stage the original bundle +and run: + +```bash +MEDMLX_TEST_CHECKPOINT_ROOT=/path/to/model-zoo \ + uv run --locked --extra dev python -m pytest -q tests/test_torch_oracles.py +``` + +The root must contain `renalStructures_CECT_segmentation/models/model.pt`. +These checks use two independent seeds for both the dense 8-cube and sparse +96-cube controls. Without this explicit asset root, those four tests are skipped; +the remaining fixture and runtime tests still run. diff --git a/src/segresnet_mlx/network.py b/src/segresnet_mlx/network.py index b2d75ac..c61917c 100644 --- a/src/segresnet_mlx/network.py +++ b/src/segresnet_mlx/network.py @@ -9,7 +9,8 @@ if TYPE_CHECKING: from ._types import MLX, Array from medmlx_core.layout import require_ncdhw -from medmlx_core.ops import conv3d_ncdhw, conv_transpose3d_ncdhw, upsample_trilinear_ncdhw +from medmlx_core.ops import conv_transpose3d_ncdhw, upsample_trilinear_ncdhw +from medmlx_core.precision import Float32Operators from segresnet_mlx.topology import SegResNetSpec @@ -35,6 +36,9 @@ class SegResNet: self.spec = spec self._weights = weights self._mx = mx + # Both convolution and normalization accumulation order affect sparse + # released-checkpoint logits; correcting only one leaves amplified drift. + self._operators = Float32Operators(mx, spec.num_groups, spec.norm_eps) def __call__(self, values: Array) -> Array: require_ncdhw(values, name="SegResNet input") @@ -89,19 +93,9 @@ class SegResNet: return hidden + values def _norm(self, values: Array, prefix: str) -> Array: - mx = self._mx - groups = self.spec.num_groups - grouped = mx.reshape(values, (values.shape[0], groups, -1)) - # A single long reduction loses precision on sparse full-volume features. - mean = group_mean(grouped, mx) - centered = grouped - mean - variance = group_mean(centered * centered, mx) - scale = mx.rsqrt(variance + self.spec.norm_eps) * mx.reshape( - self._weights[f"{prefix}.weight"], (1, groups, -1) + return self._operators.norm( + values, self._weights[f"{prefix}.weight"], self._weights[f"{prefix}.bias"] ) - bias = mx.reshape(self._weights[f"{prefix}.bias"], (1, groups, -1)) - mean * scale - shape = (values.shape[0], values.shape[1], 1, 1, 1) - return values * mx.reshape(scale, shape) + mx.reshape(bias, shape) def conv( self, @@ -112,11 +106,10 @@ class SegResNet: stride: int = 1, bias: bool = False, ) -> Array: - return conv3d_ncdhw( + return self._operators.conv( values, self._weights[f"{prefix}.weight"], self._weights.get(f"{prefix}.bias") if bias else None, padding=padding, stride=stride, - mx=self._mx, ) diff --git a/src/segresnet_mlx/transforms.py b/src/segresnet_mlx/transforms.py index 02580d3..c292384 100644 --- a/src/segresnet_mlx/transforms.py +++ b/src/segresnet_mlx/transforms.py @@ -291,9 +291,16 @@ def normalize_intensity_nonzero(data: Float32Array) -> Float32Array: def sigmoid(data: Float32Array) -> Float32Array: - """``Activations(sigmoid=True)``: ``1 / (1 + exp(-x))`` in float32.""" + """Float32 sigmoid with a stable negative tail and source threshold rounding.""" one = np.float32(1.0) - return (one / (one + np.exp(-data))).astype(np.float32, copy=False) + result = np.empty_like(data) + negative_tail = data < np.float32(-1.0) + tail = np.exp(data[negative_tail]) + result[negative_tail] = tail / (one + tail) + # Keep the original expression near zero: equivalent sigmoid formulas can + # round to opposite sides of the bundle's >= 0.5 threshold there. + result[~negative_tail] = one / (one + np.exp(-data[~negative_tail])) + return result def threshold(data: Float32Array, value: float = 0.5) -> Float32Array: diff --git a/tests/test_torch_oracles.py b/tests/test_torch_oracles.py index 6ce73a4..82f80c4 100644 --- a/tests/test_torch_oracles.py +++ b/tests/test_torch_oracles.py @@ -25,8 +25,9 @@ _RAW_LOGIT_ATOL = 1e-4 @pytest.mark.parametrize("sparse_volume", [False, True], ids=["dense", "sparse-accumulation"]) +@pytest.mark.parametrize("seed", [23, 47]) def test_renal_segresnet_checkpoint_matches_monai_raw_logits( - tmp_path: Path, sparse_volume: bool + tmp_path: Path, sparse_volume: bool, seed: int ) -> None: checkpoint_root = os.environ.get("MEDMLX_TEST_CHECKPOINT_ROOT") if checkpoint_root is None: @@ -60,7 +61,7 @@ def test_renal_segresnet_checkpoint_matches_monai_raw_logits( state = state.get("model", state) source.load_state_dict(state, strict=True) - rng = np.random.default_rng(23) + rng = np.random.default_rng(seed) if sparse_volume: # Long sparse groups exposed reduction drift missed by the small controls. size = 96 @@ -84,6 +85,6 @@ def test_renal_segresnet_checkpoint_matches_monai_raw_logits( assert host.shape == expected.shape assert host.dtype == np.float32 assert np.isfinite(host).all() - assert np.allclose(host, expected, rtol=_RAW_LOGIT_RTOL, atol=_RAW_LOGIT_ATOL) + np.testing.assert_allclose(host, expected, rtol=_RAW_LOGIT_RTOL, atol=_RAW_LOGIT_ATOL) assert np.array_equal(host >= 0, expected >= 0) diff --git a/tests/test_transforms.py b/tests/test_transforms.py index 58dd266..5957b29 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -8,6 +8,7 @@ from typing import cast import numpy as np import pytest +from _contracts import from_numpy from hypothesis import given, settings from hypothesis import strategies as st from reference_data import ( @@ -18,7 +19,7 @@ from reference_data import ( write_inputs, ) -from segresnet_mlx import medmlx +from segresnet_mlx import medmlx, transforms from segresnet_mlx._arrays import MaskArray from segresnet_mlx._nifti import nib from segresnet_mlx.medmlx import BUNDLES @@ -30,6 +31,21 @@ RENAL_KEYS = BUNDLE_INPUT_KEYS["renal"] BRATS_KEYS = BUNDLE_INPUT_KEYS["brats"] +def test_sigmoid_avoids_overflow_and_preserves_source_threshold_rounding() -> None: + limit = np.finfo(np.float32).max + values = np.array( + [-limit, -1000, -100, -90, -80, -1, -1e-7, 0, 1e-7, 1, 80, 90, 100, limit], + dtype=np.float32, + ) + expected = from_numpy(values).sigmoid().numpy() + with np.errstate(over="raise", invalid="raise"): + actual = transforms.sigmoid(values) + assert actual.dtype == np.float32 + assert np.isfinite(actual).all() + np.testing.assert_allclose(actual, expected, rtol=0, atol=1e-6) + np.testing.assert_array_equal(transforms.threshold(actual), expected >= 0.5) + + @pytest.mark.parametrize("case", CASES) def test_renal_transforms_match_the_bundle(tmp_path: Path, case: str) -> None: arrays, _ = bundle_fixture("renal")