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))