From 04f5b8464ee889b52ed9f6f496fcdef9b01af8b9 Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 20 Aug 2026 12:37:23 +0000 Subject: [PATCH 1/2] refactor: extract the tail repeated in nine apply_transform methods gpu/contrast.py ended nine of its apply_transform loops with the same three things: capture the per-sample mean/std when retain_stats is set, restore them afterwards, then run region selection and drop the channel if the result went non-finite. Written out in full every time. That is _channel_stats, _restore_stats and _select_and_check now: -173 lines, +92. The per-transform loop structure is deliberately left alone rather than inverted into a callback. Every draw in this file happens inside those loops, so keeping them means the RNG consumption order is unchanged by construction, which is what makes the claim below checkable. Verified byte-for-byte rather than by inspection: pushing a fixed volume through all 24 shipped GPU configs under a fixed seed and hashing the output gives identical digests before and after. The test suite would not have caught a reordered draw; this does. _select_and_check takes the transform rather than its three region attributes. All nine call sites read exactly self.in_seg, self.out_seg and self.mix_in_out, and spelling them out as keyword arguments made the call longer than the code it replaced -- the first attempt at this was +224/-174, which is not a simplification. The trainer tail (the same nnU-Net transform sequence in eight get_training_transforms methods) is left for the next commit: nnunetv2 is an optional extra and is not installed in CI, so that one cannot be verified the same way and deserves its own review. Co-Authored-By: Claude Opus 5 --- smauglab/transforms/gpu/contrast.py | 265 ++++++++++------------------ unit_tests/test_region_and_stats.py | 67 +++++++ 2 files changed, 159 insertions(+), 173 deletions(-) diff --git a/smauglab/transforms/gpu/contrast.py b/smauglab/transforms/gpu/contrast.py index 0c33de4..bbfa0c2 100644 --- a/smauglab/transforms/gpu/contrast.py +++ b/smauglab/transforms/gpu/contrast.py @@ -31,6 +31,53 @@ def _choose_region_mode(p_in: float, p_out: float, seg_mask: torch.Tensor | None return "all" +def _channel_stats(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Per-sample mean and std of a `[N, ...spatial]` channel, each shaped `[N]`.""" + reduce_dims = tuple(range(1, x.dim())) + return x.mean(dim=reduce_dims), x.std(dim=reduce_dims) + + +def _restore_stats(x: torch.Tensor, stats: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: + """Rescale `x` so its per-sample mean and std match `stats` again. + + What `retain_stats=True` means, written out identically in nine `apply_transform` + methods before this. + """ + orig_means, orig_stds = stats + eps = 1e-8 + reduce_dims = tuple(range(1, x.dim())) + # broadcast the stats over the spatial dims: [N, 1, 1, ...] + shape = [x.shape[0]] + [1] * (x.dim() - 1) + new_mean = x.mean(dim=reduce_dims).view(shape) + new_std = x.std(dim=reduce_dims).view(shape) + return (x - new_mean) / (new_std + eps) * orig_stds.view(shape) + orig_means.view(shape) + + +def _select_and_check( + transform: ImageOnlyTransform, + orig: torch.Tensor, + x: torch.Tensor, + seg_mask: torch.Tensor | None, + note: str = "", +) -> torch.Tensor | None: + """Apply region selection, then reject the result if it went non-finite. + + Returns None when the channel should be left as it was -- the "Final safety" check + that closed all nine of these loops identically. + + Takes the transform rather than its three region attributes: every call site read + exactly `self.in_seg`, `self.out_seg` and `self.mix_in_out`, and spelling them out + made the call longer than the code it replaces. + """ + if seg_mask is not None: + region_mode = _choose_region_mode(transform.in_seg, transform.out_seg, seg_mask) + x = _apply_region_mode(orig, x, seg_mask, region_mode, mix_in_out=transform.mix_in_out) + if torch.isnan(x).any() or torch.isinf(x).any(): + print(f"Warning nan: {type(transform).__name__}{note}", flush=True) + return None + return x + + def _foreground(mask: torch.Tensor, dim: int) -> torch.Tensor: """Which voxels the segmentation covers, reducing over the class axis. @@ -208,11 +255,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ channel_data = input[:, c] # [N, ...spatial...] orig = channel_data.clone() - if self.retain_stats: - reduce_dims = tuple(range(1, channel_data.dim())) - # store per-sample mean/std (shape [N]) - orig_means = channel_data.mean(dim=reduce_dims) - orig_stds = channel_data.std(dim=reduce_dims) + stats = _channel_stats(channel_data) if self.retain_stats else None # The asserts below restate what get_kernel guarantees per kernel_type: # only Scharr yields a list, and only its branch iterates. @@ -252,30 +295,14 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ alpha = torch.rand(1, device=input.device) x = alpha * orig + (1 - alpha) * x - if self.retain_stats: - # Adjust mean and std to match original - eps = 1e-8 - reduce_dims = tuple(range(1, x.dim())) - new_mean = x.mean(dim=reduce_dims) # [N] - new_std = x.std(dim=reduce_dims) # [N] - # reshape stats to broadcast over spatial dims: [N,1,1,...] - shape = [x.shape[0]] + [1] * (x.dim() - 1) - nm = new_mean.view(shape) - ns = new_std.view(shape) - om = orig_means.view(shape) - os = orig_stds.view(shape) - x = (x - nm) / (ns + eps) * os + om + if stats is not None: + x = _restore_stats(x, stats) # Apply region selection - if seg_mask is not None: - region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) - x = _apply_region_mode(orig, x, seg_mask, region_mode, mix_in_out=self.mix_in_out) - - # Final safety: check if nan/inf appeared - if torch.isnan(x).any() or torch.isinf(x).any(): - print(f"Warning nan: {self.__class__.__name__} with kernel={self.kernel_type}", flush=True) + checked = _select_and_check(self, orig, x, seg_mask, f" with kernel={self.kernel_type}") + if checked is None: continue - input[:, c] = x + input[:, c] = checked return input @@ -372,15 +399,10 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ orig = input[:, c] x = orig + noise - if seg_mask is not None: - region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) - x = _apply_region_mode(orig, x, seg_mask, region_mode, mix_in_out=self.mix_in_out) - - # Final safety: check if nan/inf appeared - if torch.isnan(x).any() or torch.isinf(x).any(): - print(f"Warning nan: {self.__class__.__name__}", flush=True) + checked = _select_and_check(self, orig, x, seg_mask) + if checked is None: continue - input[:, c] = x + input[:, c] = checked return input @@ -445,14 +467,10 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ x = channel_data.clone() for i in range(input.shape[0]): x[i] = x[i] * factor[i] - if seg_mask is not None: - region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) - x = _apply_region_mode(orig, x, seg_mask, region_mode, mix_in_out=self.mix_in_out) - # Final safety: check if nan/inf appeared - if torch.isnan(x).any() or torch.isinf(x).any(): - print(f"Warning nan: {self.__class__.__name__}", flush=True) + checked = _select_and_check(self, orig, x, seg_mask) + if checked is None: continue - input[:, c] = x + input[:, c] = checked return input @@ -510,11 +528,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ channel_data = -input[:, c] if self.invert_image else input[:, c] orig_full = input[:, c].clone() - if self.retain_stats: - reduce_dims = tuple(range(1, channel_data.dim())) - # store per-sample mean/std (shape [N]) - orig_means = channel_data.mean(dim=reduce_dims) - orig_stds = channel_data.std(dim=reduce_dims) + stats = _channel_stats(channel_data) if self.retain_stats else None if self.same_on_batch: gamma = ( @@ -547,30 +561,15 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ # Apply gamma transform per batch element channel_data = torch.pow(((channel_data - minm) / (rnge + 1e-8)), gamma) * rnge + minm - if self.retain_stats: - # Adjust mean and std to match original - eps = 1e-8 - reduce_dims = tuple(range(1, channel_data.dim())) - new_mean = channel_data.mean(dim=reduce_dims) # [N] - new_std = channel_data.std(dim=reduce_dims) # [N] - # reshape stats to broadcast over spatial dims: [N,1,1,...] - shape = [channel_data.shape[0]] + [1] * (channel_data.dim() - 1) - nm = new_mean.view(shape) - ns = new_std.view(shape) - om = orig_means.view(shape) - os = orig_stds.view(shape) - channel_data = (channel_data - nm) / (ns + eps) * os + om + if stats is not None: + channel_data = _restore_stats(channel_data, stats) if self.invert_image: channel_data = -channel_data - if seg_mask is not None: - region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) - channel_data = _apply_region_mode(orig_full, channel_data, seg_mask, region_mode, mix_in_out=self.mix_in_out) - # Final safety: check if nan/inf appeared - if torch.isnan(channel_data).any() or torch.isinf(channel_data).any(): - print(f"Warning nan: {self.__class__.__name__}", flush=True) + checked = _select_and_check(self, orig_full, channel_data, seg_mask) + if checked is None: continue - input[:, c] = channel_data + input[:, c] = checked return input @@ -623,11 +622,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ for c in self.apply_to_channel: channel_data = input[:, c] # [N, ...spatial...] orig = channel_data.clone() - if self.retain_stats: - reduce_dims = tuple(range(1, channel_data.dim())) - # store per-sample mean/std (shape [N]) - orig_means = channel_data.mean(dim=reduce_dims) - orig_stds = channel_data.std(dim=reduce_dims) + stats = _channel_stats(channel_data) if self.retain_stats else None if self.same_on_batch: factor = ( @@ -648,27 +643,12 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ mean = x[i].mean() x[i] = (x[i] - mean) * factor[i] + mean - if self.retain_stats: - # Adjust mean and std to match original - eps = 1e-8 - reduce_dims = tuple(range(1, x.dim())) - new_mean = x.mean(dim=reduce_dims) # [N] - new_std = x.std(dim=reduce_dims) # [N] - # reshape stats to broadcast over spatial dims: [N,1,1,...] - shape = [x.shape[0]] + [1] * (x.dim() - 1) - nm = new_mean.view(shape) - ns = new_std.view(shape) - om = orig_means.view(shape) - os = orig_stds.view(shape) - x = (x - nm) / (ns + eps) * os + om - if seg_mask is not None: - region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) - x = _apply_region_mode(orig, x, seg_mask, region_mode, mix_in_out=self.mix_in_out) - # Final safety: check if nan/inf appeared - if torch.isnan(x).any() or torch.isinf(x).any(): - print(f"Warning nan: {self.__class__.__name__}", flush=True) + if stats is not None: + x = _restore_stats(x, stats) + checked = _select_and_check(self, orig, x, seg_mask) + if checked is None: continue - input[:, c] = x + input[:, c] = checked return input @@ -721,11 +701,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ for c in self.apply_to_channel: x = input[:, c] # shape [N, ...spatial...] orig = x.clone() - if self.retain_stats: - reduce_dims = tuple(range(1, x.dim())) - # store per-sample mean/std (shape [N]) - orig_means = x.mean(dim=reduce_dims) - orig_stds = x.std(dim=reduce_dims) + stats = _channel_stats(x) if self.retain_stats else None # Normalize to make values >=0, per sample. # @@ -742,27 +718,12 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ # Apply function x = self.func(x) - if self.retain_stats: - # Adjust mean and std to match original - eps = 1e-8 - reduce_dims = tuple(range(1, x.dim())) - new_mean = x.mean(dim=reduce_dims) # [N] - new_std = x.std(dim=reduce_dims) # [N] - # reshape stats to broadcast over spatial dims: [N,1,1,...] - shape = [x.shape[0]] + [1] * (x.dim() - 1) - nm = new_mean.view(shape) - ns = new_std.view(shape) - om = orig_means.view(shape) - os = orig_stds.view(shape) - x = (x - nm) / (ns + eps) * os + om - if seg_mask is not None: - region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) - x = _apply_region_mode(orig, x, seg_mask, region_mode, mix_in_out=self.mix_in_out) - # Final safety: check if nan/inf appeared - if torch.isnan(x).any() or torch.isinf(x).any(): - print(f"Warning nan: {self.__class__.__name__}", flush=True) + if stats is not None: + x = _restore_stats(x, stats) + checked = _select_and_check(self, orig, x, seg_mask) + if checked is None: continue - input[:, c] = x + input[:, c] = checked return input @@ -896,11 +857,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ channel_data = input[:, c].clone() # shape [N, ...spatial...] orig = channel_data.clone() - if self.retain_stats: - reduce_dims = tuple(range(1, channel_data.dim())) - # store per-sample mean/std (shape [N]) - orig_means = channel_data.mean(dim=reduce_dims) - orig_stds = channel_data.std(dim=reduce_dims) + stats = _channel_stats(channel_data) if self.retain_stats else None # Process each batch element independently batch_size = channel_data.shape[0] @@ -933,28 +890,13 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ alpha = torch.rand(1, device=input.device) channel_data[b] = alpha * orig[b] + (1 - alpha) * channel_data[b] - if self.retain_stats: - # Adjust mean and std to match original - eps = 1e-8 - reduce_dims = tuple(range(1, channel_data.dim())) - new_mean = channel_data.mean(dim=reduce_dims) # [N] - new_std = channel_data.std(dim=reduce_dims) # [N] - # reshape stats to broadcast over spatial dims: [N,1,1,...] - shape = [channel_data.shape[0]] + [1] * (channel_data.dim() - 1) - nm = new_mean.view(shape) - ns = new_std.view(shape) - om = orig_means.view(shape) - os = orig_stds.view(shape) - channel_data = (channel_data - nm) / (ns + eps) * os + om + if stats is not None: + channel_data = _restore_stats(channel_data, stats) - if seg_mask is not None: - region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) - channel_data = _apply_region_mode(orig, channel_data, seg_mask, region_mode, mix_in_out=self.mix_in_out) - # Final safety: check if nan/inf appeared - if torch.isnan(channel_data).any() or torch.isinf(channel_data).any(): - print(f"Warning nan: {self.__class__.__name__}", flush=True) + checked = _select_and_check(self, orig, channel_data, seg_mask) + if checked is None: continue - input[:, c] = channel_data + input[:, c] = checked return input @@ -1133,14 +1075,10 @@ def apply_transform( nm = new_mean.view(shape) ns = new_std.view(shape) channel = (channel - nm) / (ns + eps) * os + om - if seg_mask is not None: - region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) - channel = _apply_region_mode(orig, channel, seg_mask, region_mode, mix_in_out=self.mix_in_out) - # Final safety: check if nan/inf appeared - if torch.isnan(channel).any() or torch.isinf(channel).any(): - print(f"Warning nan: {self.__class__.__name__}", flush=True) + checked = _select_and_check(self, orig, channel, seg_mask) + if checked is None: continue - input[:, c] = channel + input[:, c] = checked return input @@ -1193,11 +1131,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ for c in self.apply_to_channel: channel_data = input[:, c] # [N, ...spatial...] orig = channel_data.clone() - if self.retain_stats: - reduce_dims = tuple(range(1, channel_data.dim())) - # store per-sample mean/std (shape [N]) - orig_means = channel_data.mean(dim=reduce_dims) - orig_stds = channel_data.std(dim=reduce_dims) + stats = _channel_stats(channel_data) if self.retain_stats else None if self.same_on_batch: min_percentile = torch.rand(1, device=input.device, dtype=input.dtype) * self.max_clamp_amount @@ -1216,27 +1150,12 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ max_val = torch.quantile(x[i].flatten(), max_percentile) x[i] = torch.clamp(x[i], min_val, max_val) - if self.retain_stats: - # Adjust mean and std to match original - eps = 1e-8 - reduce_dims = tuple(range(1, x.dim())) - new_mean = x.mean(dim=reduce_dims) # [N] - new_std = x.std(dim=reduce_dims) # [N] - # reshape stats to broadcast over spatial dims: [N,1,1,...] - shape = [x.shape[0]] + [1] * (x.dim() - 1) - nm = new_mean.view(shape) - ns = new_std.view(shape) - om = orig_means.view(shape) - os = orig_stds.view(shape) - x = (x - nm) / (ns + eps) * os + om - if seg_mask is not None: - region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) - x = _apply_region_mode(orig, x, seg_mask, region_mode, mix_in_out=self.mix_in_out) - # Final safety: check if nan/inf appeared - if torch.isnan(x).any() or torch.isinf(x).any(): - print(f"Warning nan: {self.__class__.__name__}", flush=True) + if stats is not None: + x = _restore_stats(x, stats) + checked = _select_and_check(self, orig, x, seg_mask) + if checked is None: continue - input[:, c] = x + input[:, c] = checked return input diff --git a/unit_tests/test_region_and_stats.py b/unit_tests/test_region_and_stats.py index caea01c..c2c1d6e 100644 --- a/unit_tests/test_region_and_stats.py +++ b/unit_tests/test_region_and_stats.py @@ -244,3 +244,70 @@ def test_the_3d_path_is_unchanged(self): out, _ = aug_redistribute_seg(image.clone(), seg, in_seg=1.0) self.assertEqual(tuple(out.shape), tuple(image.shape)) + + +class _Stub: + """Stands in for a transform: _select_and_check only reads these three attributes.""" + + def __init__(self, in_seg: float = 0.0, out_seg: float = 0.0, mix_in_out: bool = False): + self.in_seg = in_seg + self.out_seg = out_seg + self.mix_in_out = mix_in_out + + +class TestExtractedTailHelpers(SmaugLabTestCase): + """The three helpers that replaced the tail repeated in nine apply_transform methods.""" + + def test_channel_stats_are_per_sample(self): + from smauglab.transforms.gpu.contrast import _channel_stats + + x = torch.stack([torch.full((4, 4, 4), 2.0), torch.full((4, 4, 4), 9.0)]) + + means, stds = _channel_stats(x) + + self.assertEqual(tuple(means.shape), (2,)) + self.assertAlmostEqual(float(means[0]), 2.0, places=5) + self.assertAlmostEqual(float(means[1]), 9.0, places=5) + self.assertTrue(bool((stds == 0).all())) + + def test_restore_stats_puts_mean_and_std_back(self): + from smauglab.transforms.gpu.contrast import _channel_stats, _restore_stats + + x = torch.rand(3, 5, 5, 5) * 4.0 + 1.0 + stats = _channel_stats(x) + + restored = _restore_stats(x * 100.0 - 7.0, stats) + + means, stds = _channel_stats(restored) + for b in range(3): + with self.subTest(sample=b): + self.assertAlmostEqual(float(means[b]), float(stats[0][b]), places=4) + self.assertAlmostEqual(float(stds[b]), float(stats[1][b]), places=4) + + def test_select_and_check_rejects_a_non_finite_result(self): + from smauglab.transforms.gpu.contrast import _select_and_check + + orig = torch.zeros(1, 4, 4, 4) + broken = torch.full((1, 4, 4, 4), float("nan")) + + self.assertIsNone(_select_and_check(_Stub(), orig, broken, None)) + + def test_select_and_check_passes_a_finite_result_through(self): + from smauglab.transforms.gpu.contrast import _select_and_check + + orig = torch.zeros(1, 4, 4, 4) + fine = torch.ones(1, 4, 4, 4) + + self.assertTrue(torch.equal(_select_and_check(_Stub(), orig, fine, None), fine)) + + def test_select_and_check_applies_the_region_mode(self): + from smauglab.transforms.gpu.contrast import _select_and_check + + orig = torch.zeros(1, 4, 4, 4) + transformed = torch.ones(1, 4, 4, 4) + mask = torch.zeros(1, 1, 4, 4, 4) + mask[0, 0, :2, :2, :2] = 1.0 + + out = _select_and_check(_Stub(in_seg=1.0), orig, transformed, mask) + + self.assertEqual(int(out.sum()), 8) From a368a8cb9820694625ec1783380272194e76b299 Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 27 Aug 2026 13:00:10 +0000 Subject: [PATCH 2/2] fix: type the region-selection helper's contract as a Protocol mypy on CI rejected `transform.in_seg`: the parameter was annotated as ImageOnlyTransform, and these are plain attributes on an nn.Module, so reading them through the base class resolves via Module.__getattr__ -- typed as returning `Tensor | Module`, not float. A Protocol naming in_seg, out_seg and mix_in_out fixes it and is the better annotation anyway: it states what the helper actually needs rather than naming a base class it does not otherwise use. That required ZscoreNormalizationGPU to have a mix_in_out attribute; it has the two seg knobs but no mixing. It is set in __init__ rather than added as a parameter, because the registry derives a transform's accepted config keys from its __init__ signature and a parameter would invent a setting nobody asked for. The previous `getattr(transform, "mix_in_out", False)` is gone -- mypy narrows getattr with a literal name, so it did not type-check either. Local mypy (2.1.0, older torch stubs) did not flag this; CI's does. Co-Authored-By: Claude Opus 5 --- smauglab/transforms/gpu/contrast.py | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/smauglab/transforms/gpu/contrast.py b/smauglab/transforms/gpu/contrast.py index bbfa0c2..49e0b7f 100644 --- a/smauglab/transforms/gpu/contrast.py +++ b/smauglab/transforms/gpu/contrast.py @@ -1,6 +1,6 @@ import math from collections.abc import Callable -from typing import Any, Union +from typing import Any, Protocol, Union import torch import torchvision.transforms._functional_tensor as F_t @@ -53,8 +53,26 @@ def _restore_stats(x: torch.Tensor, stats: tuple[torch.Tensor, torch.Tensor]) -> return (x - new_mean) / (new_std + eps) * orig_stds.view(shape) + orig_means.view(shape) +class _RegionSelecting(Protocol): + """What `_select_and_check` needs off the transform it is handed. + + A Protocol rather than `ImageOnlyTransform`, because these are plain attributes on + an nn.Module: reading them through the base class resolves via `Module.__getattr__`, + which is typed as returning `Tensor | Module`. Naming them here is both what makes + that type-check and a statement of the helper's actual requirement. + + All three are attributes rather than constructor parameters as far as this + Protocol is concerned; the registry derives a transform's config surface from its + __init__ signature, so declaring one here does not make it settable in a config. + """ + + in_seg: float + out_seg: float + mix_in_out: bool + + def _select_and_check( - transform: ImageOnlyTransform, + transform: _RegionSelecting, orig: torch.Tensor, x: torch.Tensor, seg_mask: torch.Tensor | None, @@ -1184,6 +1202,11 @@ def __init__( self.apply_to_channel = apply_to_channel self.in_seg = in_seg self.out_seg = out_seg + # Not a constructor parameter: this transform has no mix_in_out knob, and the + # registry derives a config's accepted keys from __init__, so adding one there + # would invent a setting. Set here so the region-selection helper's contract + # holds for every transform that uses it. + self.mix_in_out = False @torch.no_grad() def apply_transform(