From 91fc7a4718ab7ebc744f84cb5f9811392ee842b6 Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 20 Aug 2026 12:01:04 +0000 Subject: [PATCH] fix: make the random-choose bucket usable, and seed every draw with torch * RandomChooseXTransformsGPU wrote into the caller's batch: `out = input`, then `out[i:i+1] = xi` per element. It clones now, as every sibling in gpu/spatial.py already does. * The bucket could not run any transform with a kornia `_param_generator`. It calls apply_transform directly, which skips the forward_parameters step that fills `params`, so RandomLowResTransformGPU and friends raised "params must contain 'scale'" from inside a bucket. It now samples those parameters itself for children that have a generator, and leaves the contrast transforms -- which sample inside apply_transform -- on the empty-params path they expect. * RandomLowResTransformGPU read flags["data_keys"] unguarded. Only MaskSequentialOpsCustom injects that key, so every other caller got a KeyError: calling the transform standalone, or from inside the bucket, which passes the transform's own flags. It defaults to IMAGE, which is what those callers mean. * Blur sigmas and RandConv kernel sizes were drawn with Python's `random`. torch.manual_seed does not reach it, so a "seeded" training run was not reproducible, and under DDP each rank has its own `random` state and picked a different sigma for the same batch. The suite hid this because unit_tests/helpers.py::seed_everything seeds torch, numpy *and* random -- training does not call that. smauglab/transforms/rng.py is where those draws live now, built from the _next_shared_seed / _shared_rand pair that was already sitting in gpu/fromSeg.py, written for exactly this and never called once. That dead copy is deleted rather than left as a third way to draw a random number. unit_tests/test_bucket_and_rng.py: five checks fail against the previous implementation. The RNG one had to be verified separately, by restoring the random.choice call site on its own -- the other four make the module unimportable if reverted together. Models trained before this change saw the old behaviour and will not reproduce against it. No config key, parameter or default changed. Co-Authored-By: Claude Opus 5 --- smauglab/transforms/gpu/contrast.py | 4 +- smauglab/transforms/gpu/fromSeg.py | 34 +---- smauglab/transforms/gpu/spatial.py | 11 +- smauglab/transforms/gpu/transforms_list.py | 18 ++- smauglab/transforms/rng.py | 70 ++++++++++ unit_tests/test_bucket_and_rng.py | 152 +++++++++++++++++++++ 6 files changed, 250 insertions(+), 39 deletions(-) create mode 100644 smauglab/transforms/rng.py create mode 100644 unit_tests/test_bucket_and_rng.py diff --git a/smauglab/transforms/gpu/contrast.py b/smauglab/transforms/gpu/contrast.py index 358734b..d3bae6d 100644 --- a/smauglab/transforms/gpu/contrast.py +++ b/smauglab/transforms/gpu/contrast.py @@ -1,5 +1,4 @@ import math -import random from collections.abc import Callable from typing import Any, Union @@ -9,6 +8,7 @@ from torch.nn import functional as F from smauglab.transforms.gpu.base import ImageOnlyTransform +from smauglab.transforms.rng import shared_choice def _choose_region_mode(p_in: float, p_out: float, seg_mask: torch.Tensor | None) -> str: # noqa: ARG001 -- seg_mask kept for signature symmetry with _apply_region_mode @@ -216,7 +216,7 @@ def get_kernel(self, device: torch.device) -> Union[Tensor, list[Tensor]]: kernel = get_gaussian_kernel3d(kernel_size, sigma, torch.float32, device) elif self.kernel_type == "RandConv": # choose random odd kernel size e.g. [1,3,5,7] - k = int(random.choice(self.kernel_sizes)) # define kernel_sizes in __init__ + k = int(shared_choice(self.kernel_sizes)) # define kernel_sizes in __init__ std = 1.0 / math.sqrt(k * k) kernel = torch.randn((k, k, k), device=device) * std # for 3D diff --git a/smauglab/transforms/gpu/fromSeg.py b/smauglab/transforms/gpu/fromSeg.py index 5afdf9b..4ed6a75 100644 --- a/smauglab/transforms/gpu/fromSeg.py +++ b/smauglab/transforms/gpu/fromSeg.py @@ -1,12 +1,11 @@ -import random from typing import Any import torch -import torch.distributed as dist from torch import Tensor, nn from torch.nn import functional as F from smauglab.transforms.gpu.base import ImageOnlyTransform +from smauglab.transforms.rng import shared_choice # ── PALETTE AUG helpers ────────────────────────────────────────────────── @@ -432,7 +431,7 @@ def apply_transform( synth = torch.stack(synth_list) # (B, N) synth_01 = synth.reshape(B, 1, D, H, W) - sigma = random.choice(self.blur_sigmas) + sigma = shared_choice(self.blur_sigmas) if sigma > 0.0: synth_01 = _gaussian_blur_3d(synth_01, sigma) synth = synth_01.reshape(B, N) @@ -472,7 +471,7 @@ def apply_transform( # ── Step 3: optional second blur, then foreground z-score ───────────── synth_01 = synth.reshape(B, 1, D, H, W) - sigma2 = random.choice(self.blur_sigmas) + sigma2 = shared_choice(self.blur_sigmas) if sigma2 > 0.0: synth_01 = _gaussian_blur_3d(synth_01, sigma2) synth = synth_01.reshape(B, N) @@ -489,20 +488,6 @@ def apply_transform( return out -_SHARED_RNG_COUNTER = 0 - - -def _next_shared_seed() -> int: - global _SHARED_RNG_COUNTER # noqa: PLW0603 -- module-level counter is the point: it makes successive seeds distinct - _SHARED_RNG_COUNTER += 1 - seed = (int(torch.initial_seed()) + _SHARED_RNG_COUNTER) % (2**63 - 1) - if dist.is_available() and dist.is_initialized(): - seed_tensor = torch.tensor([seed], dtype=torch.long) - dist.broadcast(seed_tensor, src=0) - seed = int(seed_tensor.item()) - return seed - - def _minmax_norm(x: torch.Tensor, eps: float = 1e-8) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Per-sample min-max normalise to [0, 1]. Returns (normed, min, max).""" B = x.shape[0] @@ -532,19 +517,6 @@ def _zscore_renorm(x: torch.Tensor, bg_threshold: float = 1e-6) -> torch.Tensor: return torch.where(fg, (x - mean) / std, torch.zeros_like(x)) -def _shared_cpu_generator() -> torch.Generator: - generator = torch.Generator(device="cpu") - generator.manual_seed(_next_shared_seed()) - return generator - - -def _shared_rand(shape: tuple[int, ...], device: torch.device, dtype: torch.dtype) -> torch.Tensor: - if not (dist.is_available() and dist.is_initialized()): - return torch.rand(shape, device=device, dtype=dtype) - rand_cpu = torch.rand(shape, generator=_shared_cpu_generator(), device="cpu", dtype=dtype) - return rand_cpu.to(device=device, dtype=dtype) - - def collapse_onehot_to_index(seg_raw: torch.Tensor) -> torch.Tensor: """ Convert a one-hot segmentation mask to a single-channel integer index mask. diff --git a/smauglab/transforms/gpu/spatial.py b/smauglab/transforms/gpu/spatial.py index 6801040..932900d 100644 --- a/smauglab/transforms/gpu/spatial.py +++ b/smauglab/transforms/gpu/spatial.py @@ -232,12 +232,17 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ scales = params["scale"] # shape [B, 3] - if flags["data_keys"][0] is DataKey.IMAGE: + # Only MaskSequentialOpsCustom injects "data_keys" (see gpu/base.py), so a bare + # `flags["data_keys"]` raised KeyError for every other caller -- calling this + # transform standalone, or from inside RandomChooseXTransformsGPU, which passes + # the transform's own `flags`. Defaulting to IMAGE is what those callers mean. + data_keys = flags.get("data_keys") or [DataKey.INPUT] + if data_keys[0] in (DataKey.INPUT, DataKey.IMAGE): resample = "trilinear" - elif flags["data_keys"][0] is DataKey.MASK: + elif data_keys[0] is DataKey.MASK: resample = "nearest" else: - raise ValueError(f"Unsupported data key {flags['data_keys'][0]} for RandomLowResTransformGPU. Expected IMAGE or MASK.") + raise ValueError(f"Unsupported data key {data_keys[0]} for RandomLowResTransformGPU. Expected IMAGE or MASK.") # Define interpolation modes interp_down = resample diff --git a/smauglab/transforms/gpu/transforms_list.py b/smauglab/transforms/gpu/transforms_list.py index 23fad4f..32bd632 100644 --- a/smauglab/transforms/gpu/transforms_list.py +++ b/smauglab/transforms/gpu/transforms_list.py @@ -723,9 +723,18 @@ def _apply_mix(self, x: Tensor, seg: Tensor | None) -> Tensor: continue if not hasattr(t, "apply_transform"): raise TypeError(f"All transforms must implement apply_transform like ImageOnlyTransform. Got {type(t)}") - # Most contrast transforms perform their random sampling inside apply_transform. + # Most contrast transforms perform their random sampling inside + # apply_transform, so an empty params dict is all they need. The ones with a + # kornia `_param_generator` (the spatial transforms) read their draw out of + # `params` instead, and calling apply_transform directly skips the + # forward_parameters step that fills it -- they used to raise + # "params must contain 'scale'" from inside a bucket. Sampling here keeps + # the bucket usable for both kinds. + t_params = child_params + if getattr(t, "_param_generator", None) is not None: + t_params = {**child_params, **t.forward_parameters(x.shape)} t_flags = getattr(t, "flags", {}) - x = t.apply_transform(x, child_params, t_flags, transform=None) + x = t.apply_transform(x, t_params, t_flags, transform=None) return x @torch.no_grad() # disable gradients for efficiency @@ -736,7 +745,10 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ return self._apply_mix(input, seg) batch_size = input.shape[0] - out = input + # A clone, not `out = input`: the loop writes back through `out[i:i+1]`, so + # without it the caller's batch is modified in place. Every sibling transform + # in gpu/spatial.py clones. + out = input.clone() for i in range(batch_size): xi = out[i : i + 1] seg_i = None diff --git a/smauglab/transforms/rng.py b/smauglab/transforms/rng.py new file mode 100644 index 0000000..ad2a086 --- /dev/null +++ b/smauglab/transforms/rng.py @@ -0,0 +1,70 @@ +"""Random draws that `torch.manual_seed` actually reaches, and that DDP ranks agree on. + +Several GPU transforms reached for Python's `random.choice` to pick a blur sigma or a +kernel size, inside an `apply_transform` that was otherwise entirely `torch.rand` +driven. Two consequences: + +* `torch.manual_seed(...)` does not seed Python's `random`, so a "seeded" run was not + reproducible. The test suite hid this -- `unit_tests/helpers.py::seed_everything` + seeds torch, numpy *and* random -- but training does not call that. +* Under DistributedDataParallel each rank has its own `random` state, so ranks picked + different sigmas for the same batch. + +`gpu/fromSeg.py` already contained `_next_shared_seed` / `_shared_rand` written for +exactly this, and never called them. That machinery lives here now, with the `choice` +helper the call sites actually needed. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TypeVar + +import torch +import torch.distributed as dist + +T = TypeVar("T") + +_SHARED_RNG_COUNTER = 0 + + +def next_shared_seed() -> int: + """A seed every rank agrees on, different on each call.""" + global _SHARED_RNG_COUNTER # noqa: PLW0603 -- module-level counter is the point: it makes successive seeds distinct + _SHARED_RNG_COUNTER += 1 + seed = (int(torch.initial_seed()) + _SHARED_RNG_COUNTER) % (2**63 - 1) + if dist.is_available() and dist.is_initialized(): + seed_tensor = torch.tensor([seed], dtype=torch.long) + dist.broadcast(seed_tensor, src=0) + seed = int(seed_tensor.item()) + return seed + + +def shared_cpu_generator() -> torch.Generator: + generator = torch.Generator(device="cpu") + generator.manual_seed(next_shared_seed()) + return generator + + +def shared_rand(shape: tuple[int, ...], device: torch.device, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """Uniform [0, 1) draws; identical across ranks when running under DDP. + + Outside DDP this is just `torch.rand`, so it stays on whatever device and generator + the caller has already seeded. + """ + if not (dist.is_available() and dist.is_initialized()): + return torch.rand(shape, device=device, dtype=dtype) + rand_cpu = torch.rand(shape, generator=shared_cpu_generator(), device="cpu", dtype=dtype) + return rand_cpu.to(device=device, dtype=dtype) + + +def shared_choice(options: Sequence[T]) -> T: + """Pick one element of `options`, using torch's RNG rather than Python's. + + The drop-in replacement for `random.choice` in a transform. + """ + if len(options) == 0: + raise ValueError("cannot choose from an empty sequence") + draw = float(shared_rand((1,), torch.device("cpu")).item()) + # torch.rand is [0, 1), so the index is already in range; the clamp is belt-and-braces. + return options[min(int(draw * len(options)), len(options) - 1)] diff --git a/unit_tests/test_bucket_and_rng.py b/unit_tests/test_bucket_and_rng.py new file mode 100644 index 0000000..f4ce81c --- /dev/null +++ b/unit_tests/test_bucket_and_rng.py @@ -0,0 +1,152 @@ +"""`RandomChooseXTransformsGPU`, and the draws that `torch.manual_seed` did not reach. + +* The bucket wrote into the caller's batch, and could not run any transform with a + kornia parameter generator: calling `apply_transform` directly skips the + `forward_parameters` step that fills `params`, so those raised + "params must contain 'scale'". +* `RandomLowResTransformGPU` read `flags["data_keys"]` unguarded, which only the mask + path injects -- so it raised `KeyError` standalone and inside a bucket. +* Blur sigmas and kernel sizes were drawn with Python's `random`, which + `torch.manual_seed` does not reach and which diverges across DDP ranks. +""" + +import torch + +from smauglab.transforms.gpu.contrast import RandomConvTransformGPU +from smauglab.transforms.gpu.spatial import RandomLowResTransformGPU +from smauglab.transforms.gpu.transforms_list import RandomChooseXTransformsGPU +from smauglab.transforms.rng import shared_choice, shared_rand +from unit_tests.helpers import SmaugLabTestCase, first_output + + +class TestLowResRunsOutsideTheMaskPath(SmaugLabTestCase): + def test_it_runs_standalone(self): + """flags['data_keys'] is only injected by MaskSequentialOpsCustom.""" + torch.manual_seed(0) + transform = RandomLowResTransformGPU(p=1.0) + volume = self.tiny_volume() + + out = first_output(transform(volume)) + + self.assertIsImageLike(out, volume, "RandomLowResTransformGPU") + + def test_it_runs_inside_a_random_choose_bucket(self): + """The bucket calls apply_transform with the transform's own flags, which + carry no data_keys either.""" + torch.manual_seed(0) + bucket = RandomChooseXTransformsGPU(transforms_list=[RandomLowResTransformGPU(p=1.0)], num_transforms=1, p=1.0) + volume = self.tiny_volume() + + out = bucket.apply_transform(volume.clone(), {}, {}, transform=None) + + self.assertIsImageLike(out, volume, "RandomLowResTransformGPU in a bucket") + + def test_an_explicit_mask_key_still_selects_nearest(self): + """The branch that does exist must keep working.""" + from kornia.constants import DataKey + + torch.manual_seed(0) + transform = RandomLowResTransformGPU(p=1.0) + seg = self.tiny_seg() + params = transform.forward_parameters(seg.shape) + + out = transform.apply_transform(seg.clone(), params, {"data_keys": [DataKey.MASK]}, transform=None) + + self.assertEqual(out.shape, seg.shape) + self.assertTrue(bool(torch.isin(out, torch.tensor([0.0, 1.0])).all()), "a mask was resampled with interpolation") + + +class TestBucketDoesNotMutateItsInput(SmaugLabTestCase): + def test_the_callers_tensor_is_left_alone(self): + torch.manual_seed(0) + bucket = RandomChooseXTransformsGPU( + transforms_list=[RandomLowResTransformGPU(p=1.0)], + num_transforms=1, + p=1.0, + same_on_batch=False, + ) + volume = torch.rand(3, 1, 12, 12, 12) + before = volume.clone() + + bucket.apply_transform(volume, {}, {}, transform=None) + + self.assertTrue(torch.equal(volume, before), "RandomChooseXTransformsGPU wrote into the caller's batch") + + def test_it_still_returns_something_transformed(self): + """Cloning must not turn the bucket into a no-op.""" + torch.manual_seed(0) + bucket = RandomChooseXTransformsGPU( + transforms_list=[RandomConvTransformGPU(kernel_type="Laplace", p=1.0)], + num_transforms=1, + p=1.0, + same_on_batch=False, + ) + volume = torch.rand(2, 1, 10, 10, 10) + + out = bucket.apply_transform(volume.clone(), {}, {}, transform=None) + + self.assertFalse(torch.allclose(out, volume, atol=1e-6)) + + def test_an_empty_bucket_is_a_no_op(self): + bucket = RandomChooseXTransformsGPU(transforms_list=[], num_transforms=0, p=1.0) + volume = self.tiny_volume() + self.assertTrue(torch.equal(bucket.apply_transform(volume.clone(), {}, {}, transform=None), volume)) + + def test_the_same_on_batch_path_also_runs_a_generator_transform(self): + torch.manual_seed(0) + bucket = RandomChooseXTransformsGPU( + transforms_list=[RandomLowResTransformGPU(p=1.0)], + num_transforms=1, + p=1.0, + same_on_batch=True, + ) + volume = self.tiny_volume() + + out = bucket.apply_transform(volume.clone(), {}, {}, transform=None) + + self.assertIsImageLike(out, volume, "bucket with same_on_batch") + + +class TestTorchSeedReachesEveryDraw(SmaugLabTestCase): + """`torch.manual_seed` alone must be enough. + + The suite's own `seed_everything` seeds torch, numpy *and* Python's `random`, which + is exactly why this went unnoticed -- training does not call it. These tests seed + only torch. + """ + + def test_randconv_is_reproducible_under_torch_seed_alone(self): + outputs = [] + for _ in range(2): + torch.manual_seed(1234) + transform = RandomConvTransformGPU(kernel_type="RandConv", p=1.0, kernel_sizes=[1, 3, 5, 7]) + outputs.append(transform.apply_transform(self.tiny_volume(), {}, {}, transform=None).clone()) + + self.assertTrue(torch.equal(outputs[0], outputs[1]), "RandomConvTransformGPU drew its kernel size from an unseeded generator") + + def test_shared_choice_covers_the_whole_sequence(self): + torch.manual_seed(0) + options = (1, 3, 5, 7) + + seen = {shared_choice(options) for _ in range(200)} + + self.assertEqual(seen, set(options)) + + def test_shared_choice_is_reproducible(self): + def draw(): + torch.manual_seed(7) + return [shared_choice((1, 3, 5, 7)) for _ in range(20)] + + self.assertEqual(draw(), draw()) + + def test_shared_choice_rejects_an_empty_sequence(self): + with self.assertRaises(ValueError): + shared_choice([]) + + def test_shared_rand_is_plain_torch_rand_outside_ddp(self): + """No process group initialised, so it must stay on the caller's generator.""" + torch.manual_seed(11) + expected = torch.rand((4,)) + torch.manual_seed(11) + + self.assertTrue(torch.equal(shared_rand((4,), torch.device("cpu")), expected))