diff --git a/smauglab/transforms/gpu/fromSeg.py b/smauglab/transforms/gpu/fromSeg.py index 023f8a5..a53155a 100644 --- a/smauglab/transforms/gpu/fromSeg.py +++ b/smauglab/transforms/gpu/fromSeg.py @@ -216,8 +216,20 @@ def _region_moments(region_mask: torch.Tensor, design: torch.Tensor) -> tuple[to `E[x^2] - E[x]^2`, which is the same quantity the explicit `((x - mean) * mask) ** 2` pass computed -- the mask is binary, so squaring it changes nothing -- without materialising the centred array. + + The matmul is pinned to fp32, and must stay that way. `autocast` demotes matmuls + to fp16, which maxes at 65504, so the `1` column -- a plain voxel count -- saturates + to `inf` for any region above about 40^3 voxels. At the 128^3 training patch size + the background region always does, and then `counts=inf` gives `means=NaN`, every + per-region Gaussian contribution is NaN, the `tmax - tmin > 1e-8` guard in + `RandomRedistributeSegGPU` is False on NaN, and the transform returns its input + untouched -- an exact identity that no NaN check sees, because the NaN never + reaches the output. fp32 represents these counts exactly (128^3 is far below 2^24) + and TF32 keeps the matmul as fast as the fp16 one. Returned in fp32 deliberately: + casting back to `design.dtype` would reintroduce the overflow under autocast. """ - sums = region_mask.to(design.dtype) @ design # (R, 3) + with torch.autocast(device_type=design.device.type, enabled=False): + sums = region_mask.to(torch.float32) @ design.to(torch.float32) # (R, 3) counts = sums[:, 0].clamp_min(1) means = sums[:, 1] / counts variances = (sums[:, 2] / counts - means * means).clamp_min(0) diff --git a/unit_tests/test_redistribute_autocast.py b/unit_tests/test_redistribute_autocast.py new file mode 100644 index 0000000..cef98bc --- /dev/null +++ b/unit_tests/test_redistribute_autocast.py @@ -0,0 +1,102 @@ +"""`RandomRedistributeSegGPU` must still redistribute under `autocast`. + +`_region_moments` computes each region's voxel count, mean and variance as one +`[R, S] x [S, 3]` matmul. `autocast` demotes matmuls to fp16, whose largest finite +value is 65504, so the `1` column -- a plain voxel count -- saturates to `inf` for +any region above about 40^3 voxels. At the 128^3 patch size nnU-Net trains on, the +background region always does. + +Nothing then raises. `counts=inf` makes `means=NaN`, every per-region Gaussian +contribution is NaN, `tmax - tmin > 1e-8` is False on NaN so the `torch.where` +returns the input, and `retain_stats` maps that back to the original mean and std -- +an exact identity. The transform's own `torch.isfinite(x).all()` guard passes, +because the NaN only ever lived in `to_add`. + +So the augmentation reported itself as applied and did nothing, for every real +training run between `20360d8` (which introduced the matmul; the reduction it +replaced was fp32 under autocast and did not overflow) and this test. It sits at +`p=0.4` in the paper's "Ours" config, the highest probability of any transfer +augmentation there. + +fp16 autocast on CPU demotes the same op, so these do not need a GPU. +""" + +from __future__ import annotations + +import torch + +from smauglab.transforms.gpu.fromSeg import RandomRedistributeSegGPU, _region_moments +from unit_tests.helpers import SmaugLabTestCase + +#: Big enough to reproduce the identity, small enough to stay a unit test. Both the +#: count column and the `sum(x)` column have to overflow: with only the count at `inf` +#: the mean comes out 0 rather than NaN, which is still wrong but still moves the patch. +#: `x` is the transform's [0, 1] working space, so `sum(x)` needs roughly 2 * 65504 +#: voxels. 48^3 is not enough; 64^3 is, and matches the real failure at 128^3. +SHAPE = (1, 1, 64, 64, 64) + + +def pair() -> tuple[torch.Tensor, torch.Tensor]: + """A z-scored patch and the single-channel integer label map the trainer passes.""" + seg = torch.zeros(*SHAPE) + seg[:, :, 16:32, 16:32, 16:32] = 1.0 + seg[:, :, 36:48, 36:48, 36:48] = 2.0 + return torch.randn(*SHAPE), seg + + +class TestRegionCountsSurviveAutocast(SmaugLabTestCase): + def test_the_count_does_not_overflow_fp16(self): + values = torch.rand(100_000) + masks = torch.ones(2, 100_000, dtype=torch.bool) + design = torch.stack([torch.ones_like(values), values, values * values], dim=1) + + with torch.autocast("cpu", torch.float16): + counts, means, stds = _region_moments(masks, design) + + self.assertTrue(torch.isfinite(counts).all(), f"counts overflowed: {counts}") + self.assertEqual(float(counts[0]), 100_000.0) + self.assertAlmostEqual(float(means[0]), float(values.mean()), places=4) + self.assertAlmostEqual(float(stds[0]), float(values.std(unbiased=False)), places=4) + + def test_it_agrees_with_the_same_call_outside_autocast(self): + values = torch.rand(100_000) + masks = torch.rand(3, 100_000) > 0.5 + design = torch.stack([torch.ones_like(values), values, values * values], dim=1) + + plain = _region_moments(masks, design) + with torch.autocast("cpu", torch.float16): + amp = _region_moments(masks, design) + + for got, want in zip(amp, plain): + self.assertTrue(torch.equal(got, want)) + + +class TestRedistributeIsNotAnIdentityUnderAutocast(SmaugLabTestCase): + def test_it_moves_the_patch(self): + for seed in range(5): + with self.subTest(seed=seed): + torch.manual_seed(seed) + img, seg = pair() + transform = RandomRedistributeSegGPU(in_seg=1.0, p=1.0, p_batch=1.0) + + with torch.autocast("cpu", torch.float16): + out = transform.apply_transform(img.clone(), {"seg": seg}, {}) + + moved = (out.float() - img).abs().max().item() + self.assertTrue(torch.isfinite(out).all()) + # The perturbation is up to 2.0 in the transform's [0, 1] working space, + # so anything near zero means the identity branch fired. + self.assertGreater(moved, 1e-2, "redistribute was an identity under autocast") + + def test_autocast_does_not_change_the_result(self): + """The moments are the only thing autocast could touch here, and they are pinned.""" + img, seg = pair() + transform = RandomRedistributeSegGPU(in_seg=1.0, p=1.0, p_batch=1.0) + + torch.manual_seed(0) + plain = transform.apply_transform(img.clone(), {"seg": seg}, {}) + torch.manual_seed(0) + with torch.autocast("cpu", torch.float16): + amp = transform.apply_transform(img.clone(), {"seg": seg}, {}) + + self.assertTrue(torch.allclose(plain, amp.float(), atol=1e-5))