diff --git a/README.md b/README.md index 4755ed2..137172c 100644 --- a/README.md +++ b/README.md @@ -237,6 +237,7 @@ column -- see "Run Monai training with SmaugLab augmentations" above. Regenerate | laplace | TA | `RandomLaplaceGPU` | `LaplaceConvTransform` | | unsharp_mask | TA | `RandomUnsharpMaskGPU` | — | | rand_conv | TA | `RandomRandConvGPU` | — | +| gin | TA | `RandomGINGPU` | — | | bias_field | TA | `RandomBiasFieldGPU` | — | | inverse | TA | `RandomInverseGPU` | — | | histogram_equal | TA | `RandomHistogramEqualizationGPU` | `HistogramEqualTransform` | diff --git a/smauglab/configs/all_augmentations.json b/smauglab/configs/all_augmentations.json index 8be9b9b..377d635 100644 --- a/smauglab/configs/all_augmentations.json +++ b/smauglab/configs/all_augmentations.json @@ -293,6 +293,32 @@ "retain_stats": false, "same_on_batch": false }, + "RandomGINGPU": { + "alpha_range": [ + 0.0, + 1.0 + ], + "apply_to_channel": [ + 0 + ], + "in_seg": 0.0, + "interm_channels": 2, + "keepdim": true, + "kernel_sizes": [ + 1, + 3 + ], + "mix_in_out": false, + "n_layers": 4, + "negative_slope": 0.01, + "out_norm": "frob", + "out_seg": 0.0, + "p": 1.0, + "p_batch": 1.0, + "padding_mode": "reflect", + "retain_stats": false, + "same_on_batch": false + }, "RandomClampGPU": { "apply_to_channel": [ 0 diff --git a/smauglab/configs/transform_params_gpu.json b/smauglab/configs/transform_params_gpu.json index 05e170b..b218068 100644 --- a/smauglab/configs/transform_params_gpu.json +++ b/smauglab/configs/transform_params_gpu.json @@ -122,6 +122,9 @@ "mix_prob": 0.5, "p": 0.2 }, + "RandomGINGPU": { + "p": 0.0 + }, "RandomClampGPU": { "max_clamp_amount": 0.2, "retain_stats": true, diff --git a/smauglab/registry.py b/smauglab/registry.py index 36805bc..46b200c 100644 --- a/smauglab/registry.py +++ b/smauglab/registry.py @@ -113,6 +113,11 @@ class AugId(str, Enum): LAPLACE = "laplace" UNSHARP_MASK = "unsharp_mask" RAND_CONV = "rand_conv" + # A separate concept from RAND_CONV rather than a second implementation of it: + # RandConv is one random kernel, GIN is a cascade of them with a non-linearity + # between the layers and the output renormalised to the input's energy. Sharing the + # id would overwrite RandomRandConvGPU's matrix cell. + GIN = "gin" BIAS_FIELD = "bias_field" INVERSE = "inverse" HISTOGRAM_EQUAL = "histogram_equal" @@ -213,6 +218,9 @@ def __init__(self, source: str, problems: list[str]) -> None: "RandomLaplaceGPU", "RandomUnsharpMaskGPU", # UnsharpMaskTransform "RandomRandConvGPU", # RandomConvTransform + # No ladder key: GIN postdates the ladder. Placed next to the single-layer + # RandConv it generalises, so a config reads the two together. + "RandomGINGPU", "RandomClampGPU", # ClampTransform "RandomGaussianNoiseGPU", # GaussianNoiseTransform "RandomGaussianBlurGPU", # GaussianBlurTransform diff --git a/smauglab/transforms/__init__.py b/smauglab/transforms/__init__.py index f3a9ece..1e0b5f2 100644 --- a/smauglab/transforms/__init__.py +++ b/smauglab/transforms/__init__.py @@ -19,6 +19,7 @@ class carries an `@register(...)` decorator, so the registry is complete once th from smauglab.transforms.gpu import contrast as gpu_contrast # noqa: F401 from smauglab.transforms.gpu import domain_transfer # noqa: F401 from smauglab.transforms.gpu import fromSeg as gpu_fromSeg +from smauglab.transforms.gpu import gin as gpu_gin # noqa: F401 from smauglab.transforms.gpu import palette as gpu_palette # noqa: F401 from smauglab.transforms.gpu import spatial as gpu_spatial # noqa: F401 from smauglab.transforms.synthseg import transforms as synthseg_transforms # noqa: F401 diff --git a/smauglab/transforms/gpu/gin.py b/smauglab/transforms/gpu/gin.py new file mode 100644 index 0000000..4fe8d5f --- /dev/null +++ b/smauglab/transforms/gpu/gin.py @@ -0,0 +1,340 @@ +"""GIN: Global Intensity Non-linear augmentation. + +An image pushed through a shallow convolutional network whose weights are drawn afresh on +every call, blended back towards the original by a per-sample `alpha`, and rescaled so the +result carries the same energy as the input. Because the network is random rather than +trained, the family of intensity mappings it realises is far wider than any fixed curve -- +which is the point: it is a single-source domain-generalisation augmentation, meant to make +a model indifferent to the appearance of the scanner it was trained on. + +Based on: + Ouyang, C., Chen, C., Li, S., Li, Z., Qin, C., Bai, W., & Rueckert, D. (2022). + Causality-inspired single-source domain generalization for medical image segmentation. + IEEE Transactions on Medical Imaging, 42(4), 1095-1106. DOI 10.1109/TMI.2022.3224067 + +Ported from the authors' 3-D implementation (`models/imagefilter3d.py` in +https://github.com/cheng-01037/Causality-Medical-Image-Domain-Generalization) and the +cleaner nnU-Net-oriented restatement of it in `dg_tta/gin.py` +(https://github.com/multimodallearning/DG-TTA). Both are MIT licensed. + +Two deliberate departures from upstream, both argued at their call sites: + +* **Reflect padding** rather than zeros, because a zero-padded 3x3x3 convolution darkens + the patch border and an nnU-Net patch is an interior crop, not a whole image. + `padding_mode="zeros"` restores upstream's behaviour exactly. +* **float32 as a floor**, because the Frobenius renormalisation overflows float16 often + enough to matter and does so silently. A float64 input keeps its precision. See the + comment in `apply_transform`. +""" + +import contextlib +from collections.abc import Sequence +from typing import Any + +import torch +from torch import Tensor +from torch.nn import functional as F + +from smauglab.registry import AugId, AugType, Backend, register +from smauglab.transforms.gpu.base import ImageOnlyTransform, autocast_active, segmentation_from +from smauglab.transforms.gpu.contrast import _channel_stats, _check_channel, _restore_stats, _select_and_check +from smauglab.transforms.rng import shared_rand + +#: Upstream's guard against dividing by a vanishing norm, kept at its published value. +#: Also what makes an all-zero patch come out all-zero instead of NaN. +_NORM_EPS = 1e-5 + +#: The two output normalisations, and the two paddings. Spelled out so a config typo is +#: rejected at construction rather than silently falling through to the other branch. +_OUT_NORMS = ("frob", "none") +_PADDING_MODES = ("reflect", "zeros") + + +def _draw_kernel_sizes(kernel_sizes: Sequence[int], n_layers: int) -> list[int]: + """One kernel size per layer, drawn from torch's RNG in a single call. + + `shared_choice` is the obvious helper and the one `RandomRandConvGPU` uses, but it + ends in `float(...item())` -- a read back to the host -- so calling it once per layer + costs `n_layers` stalls per forward. Bucketing one `shared_rand` draw costs one, and + keeps the draw somewhere `torch.manual_seed` reaches and every DDP rank agrees on. + """ + draw = shared_rand((n_layers,), torch.device("cpu")) + # `shared_rand` is [0, 1), so the index is already in range; the clamp is belt and braces. + index = (draw * len(kernel_sizes)).long().clamp_(max=len(kernel_sizes) - 1) + return [int(kernel_sizes[i]) for i in index.tolist()] + + +def _gin_stack( + work: Tensor, + widths: Sequence[int], + kernel_sizes: Sequence[int], + negative_slope: float, + padding_mode: str, + same_on_batch: bool, +) -> Tensor: + """Push a `[B, C, D, H, W]` volume through one freshly drawn conv cascade. + + `widths` is the channel count at every layer boundary, so it holds `n_layers + 1` + entries and starts and ends at `C`; `kernel_sizes` holds the one size already drawn + for each layer, not the set they were drawn from. Weights and biases are raw `N(0, 1)` + with no fan-in scaling -- the paper's `N(0, I)`, and safe only because the caller + renormalises the output; see the class docstring. Everything is drawn in `work.dtype`, + which the caller has already floored at float32. + + The batch is folded into the group axis: with the volume viewed as + `[1, B*Cin, D, H, W]`, a weight of `[B*Cout, Cin, k, k, k]` and `groups=B`, group `b` + sees channels `[b*Cin : (b+1)*Cin]` and nothing else, so every sample gets its own + independent network out of a single convolution. + + `same_on_batch` repeats one drawn network across the groups rather than switching to + `groups=1`. Two reasons: the rows then come out bitwise identical, which is what the + flag promises; and `groups=1` is the slower route for the shapes this actually runs. + Measured on an idle A40 over a `[2, 1, 128, 128, 128]` batch with `k=3`, grouped + against shared: `1 -> 2` channels 0.32 ms against 1.10 ms, `2 -> 2` 1.67 against 1.30, + `2 -> 1` 1.69 against 7.33 -- so about 5.4 ms against 11.0 ms over the default + cascade's four layers. `AugmentationSequentialCustom` forces `same_on_batch=True` onto + every child, so that is the shipped path and not the exotic one. + + `kernels.depthwise_conv3d` cannot serve here: it is depthwise by construction + (`groups = B*C`, one input channel per filter), and the interior layers map + `interm -> interm` within a group. + """ + batch = work.shape[0] + spatial = work.shape[2:] + x = work.reshape(1, batch * widths[0], *spatial) + last = len(kernel_sizes) - 1 + + for layer, k in enumerate(kernel_sizes): + channels_in, channels_out = widths[layer], widths[layer + 1] + if same_on_batch: + weight = torch.randn((channels_out, channels_in, k, k, k), device=work.device, dtype=work.dtype) + weight = weight.repeat(batch, 1, 1, 1, 1) + bias = torch.randn((channels_out,), device=work.device, dtype=work.dtype).repeat(batch) + else: + weight = torch.randn((batch * channels_out, channels_in, k, k, k), device=work.device, dtype=work.dtype) + bias = torch.randn((batch * channels_out,), device=work.device, dtype=work.dtype) + + if k > 1: + # Skipped entirely for k == 1, where upstream's `padding=k//2` is 0 anyway and + # the call would be a full-volume copy for nothing -- four of them per forward + # at the default `kernel_sizes`. + pad = [k // 2] * 6 + x = F.pad(x, pad, mode="reflect") if padding_mode == "reflect" else F.pad(x, pad) + x = F.conv3d(x, weight, groups=batch) + # One scalar per output plane, which is upstream's `[out_channel*nb, 1, 1, 1]` + # written as the view it broadcasts to rather than left to right-alignment. + x = x + bias.view(1, -1, 1, 1, 1) + if layer != last: + x = F.leaky_relu(x, negative_slope) + + return x.reshape(batch, widths[-1], *spatial) + + +@register( + aug_id=AugId.GIN, + backend=Backend.GPU, + group=AugType.TA, +) +class RandomGINGPU(ImageOnlyTransform): + """GIN: a cascade of randomly drawn convolutions, renormalised to the input's energy. + + If the image is a torch Tensor, it is expected to have [N, C, X, Y, Z] shape. + + The network is redrawn on every call, so there is no state and nothing to seed beyond + torch's own RNG. Each layer's kernel size is drawn from `kernel_sizes`, its weights and + bias from `N(0, 1)` with no fan-in scaling, and `leaky_relu` follows every layer but + the last. The result is blended with the original at a per-sample `alpha` and then + rescaled so its Frobenius norm matches the input's. + + That final rescaling is what makes the unscaled weights safe, and makes the output + scale *exactly* independent of `n_layers`, `kernel_sizes` and `interm_channels`: the + per-sample output RMS matches the input's to within 4.1e-07 relative for every + combination of those. `RandomRandConvGPU` needs its `1/sqrt(k**3)` kernel scaling precisely because it + has no such renormalisation. Note the invariant is on RMS and not on the standard + deviation -- the biases and the leaky_relu move energy into the DC term, so the + standard deviation still varies by a factor of about 1.3 across those settings. + + There is no `mix_prob` here, unlike its siblings: `alpha_range` already *is* the + probability-weighted blend back towards the original, and a `mix_prob` on top would + blend twice. + + Args: + n_layers (int): Convolution layers in the cascade. Default is 4, the paper's value; + the ablation there finds one or two too weak to imitate a domain shift and many + more unrealistically aggressive. `n_layers=1` degenerates to a single linear + `C -> C` convolution and ignores `interm_channels`. + interm_channels (int): Width of the hidden layers. Default is 2, the paper's value. + kernel_sizes (list of int): Odd kernel sizes to draw from, one draw per layer. + Default is (1, 3); upstream warns against going large with several layers, and + a size above the thinnest patch axis will make reflect padding raise. + negative_slope (float): leaky_relu slope between layers. Default is 0.01, which is + what upstream's bare `F.leaky_relu` uses. + alpha_range (tuple of float): Range the blend weight is drawn from, where 0 keeps + the original and 1 takes the network's output whole. Default is (0.0, 1.0), + the paper's U(0, 1). + out_norm (str): "frob" to rescale the output to the input's Frobenius norm, "none" + to leave it alone. Default is "frob". Leaving it off makes the output scale + track the configuration, so pair it with `retain_stats`. + padding_mode (str): "reflect" or "zeros". Default is "reflect"; see the module + docstring for why that differs from upstream. + apply_to_channel (list of int): Channel indices the network maps. They are mapped + *jointly* -- one network with `len(apply_to_channel)` inputs and outputs, so it + mixes across channels as the paper intends, and the energy match is over the + selected channels together rather than one at a time. Identical to a + per-channel loop for single-channel data. Default is [0]. + retain_stats (bool): If True, restore the original per-sample mean and std + afterwards. Default is False. + same_on_batch (bool): Draw one network, one kernel size per layer and one alpha for + the whole batch. Default is False. + in_seg (float): Probability of restricting the result to inside the segmentation. + out_seg (float): Probability of restricting the result to outside it. + mix_in_out (bool): Apply to a random subset of the mask's classes rather than all. + p (float): Probability of applying the transform. Default is 1.0. + p_batch (float): Probability of applying the transform to the batch. Default is 1.0. + keepdim (bool): Whether to keep the number of dimensions. Default is True. + + Returns: + Tensor: The image with the selected channels replaced by the GIN output. + + """ + + def __init__( + self, + n_layers: int = 4, + interm_channels: int = 2, + kernel_sizes: Sequence[int] = (1, 3), + negative_slope: float = 0.01, + alpha_range: tuple[float, float] = (0.0, 1.0), + out_norm: str = "frob", + padding_mode: str = "reflect", + apply_to_channel: Sequence[int] = (0,), + retain_stats: bool = False, + same_on_batch: bool = False, + in_seg: float = 0.0, + out_seg: float = 0.0, + mix_in_out: bool = False, + p: float = 1.0, + p_batch: float = 1.0, + keepdim: bool = True, + ) -> None: + super().__init__(p=p, p_batch=p_batch, same_on_batch=same_on_batch, keepdim=keepdim) + if n_layers < 1: + raise ValueError(f"n_layers must be at least 1, got {n_layers}. A zero-layer cascade is the identity.") + if interm_channels < 1: + raise ValueError(f"interm_channels must be at least 1, got {interm_channels}.") + if len(kernel_sizes) == 0: + raise ValueError("kernel_sizes must name at least one size to draw from.") + for k in kernel_sizes: + if k < 1 or k % 2 == 0: + raise ValueError( + f"kernel_sizes must be odd and positive, got {list(kernel_sizes)}. An even size cannot keep the patch shape." + ) + if out_norm not in _OUT_NORMS: + raise ValueError(f"out_norm must be one of {_OUT_NORMS}, got {out_norm!r}.") + if padding_mode not in _PADDING_MODES: + raise ValueError(f"padding_mode must be one of {_PADDING_MODES}, got {padding_mode!r}.") + if alpha_range[0] > alpha_range[1]: + raise ValueError(f"alpha_range must be (low, high), got {tuple(alpha_range)}.") + # Harmless on every sibling, where a repeated channel is just a second application. + # Here it changes the network: `len(apply_to_channel)` is the cascade's input width, + # so (0, 0) builds a two-input net and then scatters both of its outputs into + # channel 0, last write winning. + if len(set(apply_to_channel)) != len(list(apply_to_channel)): + raise ValueError( + f"apply_to_channel must not repeat a channel, got {list(apply_to_channel)}. " + "The channels are mapped jointly, so a duplicate changes the network's width." + ) + self.n_layers = n_layers + self.interm_channels = interm_channels + self.kernel_sizes = kernel_sizes + self.negative_slope = negative_slope + self.alpha_range = alpha_range + self.out_norm = out_norm + self.padding_mode = padding_mode + self.apply_to_channel = apply_to_channel + self.retain_stats = retain_stats + self.in_seg = in_seg + self.out_seg = out_seg + self.mix_in_out = mix_in_out + + @torch.no_grad() # disable gradients for efficiency + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: + # A clone, not the caller's tensor: this writes channels back with + # `input[:, c] = ...` and kornia passes the caller's own tensor through when + # every sample applies. + input = input.clone() + seg_mask = segmentation_from(params) + channels = [int(c) for c in self.apply_to_channel] + for c in channels: + # Before the gather, so a config typo names the key it came from rather than + # raising a bare IndexError out of advanced indexing. + _check_channel(self, c, input.shape[1]) + batch = input.shape[0] + + # float32 as a floor, for a failure that is otherwise invisible. The trainer runs + # this pipeline inside `autocast`, so `F.conv3d` would hand back float16; the + # kernels are raw N(0, 1) with no fan-in scaling, and `||mixed||_F` then passed + # float16's 65504 in 23 of 150 draws at these defaults on a 1x128^3 patch, 48 of + # 150 on 2x192^3, and 105 of 150 with kernel_sizes=(1, 3, 5, 7). The norm comes + # back `inf`, `1 / (inf + eps)` is exactly 0, and the result is an all-zero patch + # that `torch.isfinite(...).all()` reports as fine -- so neither + # `_select_and_check` below nor the trainer's non-finite-loss guard ever sees it, + # and the network trains on a blank volume. + # + # Guarded on `autocast_active()` rather than entered unconditionally: `autocast` + # on mps complains that it is not implemented even with `enabled=False`, which is + # the same reason train_step wraps its own region in a conditional. + guard = torch.autocast(input.device.type, enabled=False) if autocast_active() else contextlib.nullcontext() + with guard: + # Advanced indexing already copies; the promotion is the other half of the + # forcing, for a caller that hands over float16 storage. `promote_types` + # rather than `.float()` so the floor is float32 and not the ceiling -- a + # float64 input would otherwise be quietly halved in precision. + work = input[:, channels].to(torch.promote_types(input.dtype, torch.float32)) + widths = [len(channels), *[self.interm_channels] * (self.n_layers - 1), len(channels)] + net = _gin_stack( + work, + widths=widths, + kernel_sizes=_draw_kernel_sizes(self.kernel_sizes, self.n_layers), + negative_slope=self.negative_slope, + padding_mode=self.padding_mode, + same_on_batch=self.same_on_batch, + ) + + # One alpha per sample, shared across the selected channels. A single draw + # under `same_on_batch`, which then broadcasts over the batch axis. + low, high = self.alpha_range + draws = torch.rand(1 if self.same_on_batch else batch, device=work.device, dtype=work.dtype) + alpha = (draws * (high - low) + low).reshape(-1, 1, 1, 1, 1) + mixed = alpha * net + (1.0 - alpha) * work + + if self.out_norm == "frob": + # Upstream's `torch.norm(x.reshape(B, C, -1), dim=(-1, -2), p="fro")` + # reduces over the channel axis as well as the spatial ones, so this is + # one scalar per sample rather than one per channel -- deliberately, and + # bitwise the same value as the reshape below. Channels therefore share a + # scale factor, which couples modalities held at different scales when + # `apply_to_channel` names more than one. + energy_in = work.reshape(batch, -1).norm(dim=1).reshape(-1, 1, 1, 1, 1) + energy_out = mixed.reshape(batch, -1).norm(dim=1).reshape(-1, 1, 1, 1, 1) + # The factor is formed first. Upstream writes `mixed * (1 / (out + eps)) * + # in`, which walks the volume twice and allocates a second copy of it; the + # scale is one number per sample, so it costs nothing to combine. + mixed = mixed * (energy_in / (energy_out + _NORM_EPS)) + + out = mixed.to(input.dtype) + for index, c in enumerate(channels): + orig = input[:, c] + x = out[:, index] + if self.retain_stats: + x = _restore_stats(x, _channel_stats(orig)) + # A channel rejected here reverts on its own while the others keep the GIN + # output, which gives up the joint energy match for that one call. That is + # what every sibling does with a rejected channel, and a NaN volume is the + # worse of the two outcomes. + checked = _select_and_check(self, orig, x, seg_mask, " (gin)") + if checked is None: + continue + input[:, c] = checked + + return input diff --git a/unit_tests/test_gin.py b/unit_tests/test_gin.py new file mode 100644 index 0000000..d256011 --- /dev/null +++ b/unit_tests/test_gin.py @@ -0,0 +1,323 @@ +"""GIN, the augmentation and the two things about it that are easy to get wrong. + +`RandomGINGPU` draws a fresh convolutional network on every call from raw `N(0, 1)` +weights, with no fan-in scaling. Two consequences carry most of the tests here. + +* **The Frobenius renormalisation is what makes that safe**, and it is exact rather than + statistical: the per-sample output RMS equals the input's for every depth and kernel + size, measured to 4.1e-07. `RandomRandConvGPU` needs its `1/sqrt(k**3)` because it has + no such step, and `test_randconv_gain.py` pins that. The invariant here is on RMS, not + on the standard deviation: the biases and the leaky_relu move energy into the DC term, + so the standard deviation still varies across settings. +* **The same unscaled weights overflow float16**, and the failure is silent. Simulating + the stack without `apply_transform`'s float32 forcing, on a 32^3 patch with + `n_layers=6, kernel_sizes=(3, 5)`, the norm exceeds float16's 65504 in **19 of 20** + draws; `1 / (inf + 1e-5)` is then exactly 0 and the output is an all-zero patch that + `torch.isfinite(...).all()` reports as fine. So neither `_select_and_check` nor the + trainer's non-finite-loss guard can see it, and asserting finiteness alone would not + catch a regression. `test_a_float16_input_stays_finite_and_energy_matched` asserts the + energy match instead. + +Ouyang, C., et al. (2022). Causality-inspired single-source domain generalization for +medical image segmentation. IEEE TMI 42(4), 1095-1106. DOI 10.1109/TMI.2022.3224067 +""" + +import torch + +from smauglab.transforms.gpu.gin import RandomGINGPU +from unit_tests.helpers import SmaugLabTestCase + +#: Depth and kernel-size combinations the scale invariance has to hold across. The last +#: two are deliberately far outside the paper's defaults, where an unnormalised cascade +#: amplifies by four orders of magnitude. +CONFIGS = ( + (1, (1,)), + (1, (3,)), + (2, (1, 3)), + (4, (1, 3)), + (8, (1, 3)), + (4, (1, 3, 5, 7)), + (4, (7,)), +) + +DRAWS = 5 + + +def rms(volume: torch.Tensor) -> torch.Tensor: + """Per-sample root mean square, over channels and space together. + + The quantity GIN's Frobenius renormalisation preserves -- `||x||_F` divided by the + square root of the per-sample element count, which is fixed within a comparison. + """ + return volume.reshape(volume.shape[0], -1).pow(2).mean(dim=1).sqrt() + + +class GinTestCase(SmaugLabTestCase): + """Shared driver. `apply_transform` is called directly, as the sibling tests do.""" + + def apply(self, transform: RandomGINGPU, volume: torch.Tensor, seed: int = 0, seg: torch.Tensor | None = None) -> torch.Tensor: + torch.manual_seed(seed) + params = {} if seg is None else {"seg": seg} + return transform.apply_transform(volume.clone(), params, transform.flags) + + def batch_of_identical_rows(self, rows: int = 3) -> torch.Tensor: + """A batch whose samples are byte-identical, so any difference in the output is the draw.""" + return torch.rand(1, 1, 16, 16, 16).repeat(rows, 1, 1, 1, 1) + + +class TestTheOutputScaleIsConfigurationIndependent(GinTestCase): + def test_the_output_rms_matches_the_input_rms_for_every_depth_and_kernel_size(self): + """`out_norm="frob"` pins the output energy to the input's, whatever the cascade. + + This is why the kernels are drawn from raw `N(0, 1)` rather than scaled by + `1/sqrt(k**3)` the way `_RandomConvBaseGPU.get_kernel` scales RandConv's. + """ + volume = self.tiny_volume() + reference = rms(volume) + for n_layers, kernel_sizes in CONFIGS: + transform = RandomGINGPU(p=1.0, n_layers=n_layers, kernel_sizes=kernel_sizes) + for seed in range(DRAWS): + with self.subTest(n_layers=n_layers, kernel_sizes=kernel_sizes, seed=seed): + error = float(((rms(self.apply(transform, volume, seed)) - reference) / reference).abs().max()) + self.assertLess(error, 1e-4, f"output RMS drifted by {error:.2e} relative") + + def test_without_the_frobenius_norm_the_scale_does_track_the_configuration(self): + """The control for the test above: it measures the renormalisation, not a tautology. + + Turning it off makes the output scale a function of the cascade -- measured spread + 19197x across `CONFIGS` -- which is exactly what an unscaled random net does and + what the renormalisation exists to undo. + """ + volume = self.tiny_volume() + scales = [] + for n_layers, kernel_sizes in CONFIGS: + transform = RandomGINGPU(p=1.0, n_layers=n_layers, kernel_sizes=kernel_sizes, out_norm="none") + draws = [float(rms(self.apply(transform, volume, seed))[0]) for seed in range(DRAWS)] + scales.append(sum(d * d for d in draws) ** 0.5) + + spread = max(scales) / min(scales) + self.assertGreater(spread, 10.0, f"out_norm='none' should let the scale run away, spread was only {spread:.1f}x") + + +class TestTheDrawIsSharedExactlyWhenAsked(GinTestCase): + def test_same_on_batch_true_shares_every_draw(self): + """One network, one kernel size per layer and one alpha for the whole batch. + + Asserted bitwise: `same_on_batch` repeats the drawn weights across the groups + rather than switching to `groups=1`, so the rows are not merely close. + """ + rows = self.batch_of_identical_rows() + out = self.apply(RandomGINGPU(p=1.0, same_on_batch=True), rows) + for row in range(1, out.shape[0]): + with self.subTest(row=row): + self.assertTrue(bool(torch.equal(out[0], out[row])), "same_on_batch=True gave a sample its own draw") + + def test_same_on_batch_false_gives_every_sample_its_own_network(self): + """The `groups=B` fold really does hand group `b` its own weights. + + Nothing else in the suite reaches this branch: `test_transforms_gpu.py` and its + relatives all build `AugmentationSequentialCustom(..., same_on_batch=True)`. + """ + rows = self.batch_of_identical_rows() + out = self.apply(RandomGINGPU(p=1.0, same_on_batch=False), rows) + for row in range(1, out.shape[0]): + with self.subTest(row=row): + self.assertFalse(bool(torch.allclose(out[0], out[row])), "identical samples came out identical without same_on_batch") + + def test_a_fixed_seed_reproduces_the_output(self): + """Every draw goes through torch's RNG, including the per-layer kernel size. + + `_draw_kernel_sizes` uses `smauglab.transforms.rng`, not `random.choice`, so + `torch.manual_seed` reaches it and DDP ranks agree on it. + """ + volume = self.tiny_volume() + transform = RandomGINGPU(p=1.0) + self.assertTrue(bool(torch.equal(self.apply(transform, volume, 7), self.apply(transform, volume, 7)))) + + +class TestTheBlend(GinTestCase): + def test_alpha_range_zero_is_the_identity(self): + """`alpha=0` keeps the original, up to the renormalisation's own epsilon. + + Not bitwise: with `mixed == work` the rescaling still multiplies by + `||x|| / (||x|| + 1e-5)`, which is 1.19e-07 away from 1 on this patch. Asserted on + a random volume rather than `constant_volume()`, where that factor happens to + round to exactly 1.0 and a `torch.equal` would pass by luck. + """ + volume = self.tiny_volume() + out = self.apply(RandomGINGPU(p=1.0, alpha_range=(0.0, 0.0)), volume) + self.assertTrue(bool(torch.allclose(out, volume, atol=1e-6)), f"max deviation {float((out - volume).abs().max()):.2e}") + + def test_alpha_range_one_never_keeps_the_original(self): + """`alpha=1` takes the network's output whole, so the blend is not silently clamped.""" + volume = self.tiny_volume() + out = self.apply(RandomGINGPU(p=1.0, alpha_range=(1.0, 1.0)), volume) + self.assertGreater(float((out - volume).abs().max()), 0.1) + + +class TestPrecision(GinTestCase): + def test_a_float16_input_stays_finite_and_energy_matched(self): + """float32 is forced inside the transform, so the float16 norm cannot overflow. + + The configuration is the one measured to overflow in 19 of 20 draws without the + forcing. Energy is asserted, not just finiteness: the overflow produces an + all-zero patch that `isfinite` calls healthy. + """ + volume = torch.rand(1, 1, 32, 32, 32, dtype=torch.float16) + reference = float(rms(volume.float())[0]) + transform = RandomGINGPU(p=1.0, n_layers=6, kernel_sizes=(3, 5)) + for seed in range(DRAWS): + with self.subTest(seed=seed): + out = self.apply(transform, volume, seed) + self.assertEqual(out.dtype, torch.float16, "the input dtype was not restored") + self.assertTrue(bool(torch.isfinite(out).all()), "non-finite output") + self.assertTrue(bool(out.any()), "the output was zeroed, which is what a float16 norm overflow looks like") + got = float(rms(out.float())[0]) + self.assertAlmostEqual(got / reference, 1.0, delta=1e-3, msg=f"energy drifted to {got:.4f} from {reference:.4f}") + + def test_an_open_autocast_region_does_not_change_the_result(self): + """The `autocast_active()`-guarded `enabled=False` region keeps the cascade in float32. + + Bitwise, because the guard means the same kernels run the same convolutions + whether or not a caller has autocast open -- which is how the trainer calls it. + """ + volume = self.tiny_volume() + transform = RandomGINGPU(p=1.0) + plain = self.apply(transform, volume, 3) + with torch.autocast("cpu", dtype=torch.bfloat16, enabled=True): + under_autocast = self.apply(transform, volume, 3) + self.assertTrue(bool(torch.equal(plain, under_autocast)), "autocast changed the result, so the cascade ran in reduced precision") + + +class TestDegeneratePatches(GinTestCase): + def test_a_constant_patch_stays_constant(self): + """GIN cannot manufacture structure out of an air patch. + + A constant input gives a constant output whose magnitude the renormalisation + restores to the input's. The sign may flip -- the network's output can be + negative where the input was not -- so only the magnitude is pinned. + """ + volume = self.constant_volume() + expected = abs(float(volume.flatten()[0])) + transform = RandomGINGPU(p=1.0) + for seed in range(DRAWS): + with self.subTest(seed=seed): + out = self.apply(transform, volume, seed) + unique = torch.unique(out) + self.assertEqual(unique.numel(), 1, f"a constant patch came out with {unique.numel()} values") + self.assertAlmostEqual(abs(float(unique[0])), expected, delta=1e-3) + + def test_an_all_zero_patch_stays_all_zero(self): + """`||x||_F = 0` makes the rescaling return zero rather than NaN. + + The network's biases give it a non-zero output even here, so this is the epsilon's + placement being right: the zero numerator wins over the non-zero denominator. + """ + volume = torch.zeros_like(self.tiny_volume()) + transform = RandomGINGPU(p=1.0) + for seed in range(DRAWS): + with self.subTest(seed=seed): + out = self.apply(transform, volume, seed) + self.assertTrue(bool((out == 0).all()), "an empty patch came back non-zero") + + +class TestChannels(GinTestCase): + def test_multi_channel_input_is_transformed_jointly(self): + """One network over all of `apply_to_channel`, with a single joint energy match. + + The selected channels mix -- which is the published method -- and share one + rescaling factor, so their *joint* RMS is what is preserved, not each channel's + own. An unselected channel is left byte-identical. + """ + volume = torch.rand(2, 3, 16, 16, 16) + out = self.apply(RandomGINGPU(p=1.0, apply_to_channel=(0, 1)), volume) + + self.assertTrue(bool(torch.equal(out[:, 2], volume[:, 2])), "an unselected channel was modified") + for channel in (0, 1): + with self.subTest(channel=channel): + self.assertFalse(bool(torch.equal(out[:, channel], volume[:, channel])), "a selected channel was left alone") + + before, after = rms(volume[:, :2]), rms(out[:, :2]) + self.assertTrue(bool(torch.allclose(after, before, rtol=1e-4)), f"joint RMS moved from {before.tolist()} to {after.tolist()}") + + def test_a_duplicate_apply_to_channel_is_rejected(self): + """A repeated channel is harmless on the siblings and changes the network here. + + `len(apply_to_channel)` is the cascade's input width, so `(0, 0)` would build a + two-input net and scatter both its outputs into channel 0. + """ + with self.assertRaises(ValueError): + RandomGINGPU(p=1.0, apply_to_channel=(0, 0)) + + def test_an_out_of_range_channel_names_the_config_key(self): + """Checked before the gather, so the message names the key and not a stray index.""" + transform = RandomGINGPU(p=1.0, apply_to_channel=(1,)) + with self.assertRaises(IndexError) as caught: + self.apply(transform, self.tiny_volume()) + self.assertIn("apply_to_channel", str(caught.exception)) + + +class TestConstructorValidation(GinTestCase): + def test_invalid_constructor_arguments_are_rejected(self): + """A config typo fails at construction rather than falling through to a default.""" + cases = { + "n_layers below one": {"n_layers": 0}, + "interm_channels below one": {"interm_channels": 0}, + "no kernel sizes to draw from": {"kernel_sizes": ()}, + "an even kernel size": {"kernel_sizes": (2,)}, + "a non-positive kernel size": {"kernel_sizes": (0,)}, + "an unknown out_norm": {"out_norm": "nope"}, + "an unknown padding_mode": {"padding_mode": "nope"}, + "an inverted alpha_range": {"alpha_range": (1.0, 0.0)}, + } + for label, kwargs in cases.items(): + with self.subTest(case=label), self.assertRaises(ValueError): + RandomGINGPU(p=1.0, **kwargs) + + +class TestPadding(GinTestCase): + def test_zero_padding_changes_only_the_border(self): + """`padding_mode` is wired through, and reaches exactly as far as the cascade does. + + Four layers of `k=3` propagate the padded edge one voxel per layer, so the two + modes differ over the outer four voxels and leave everything inside them + byte-identical. Measured on the network's raw output (`alpha=1`, `out_norm="none"`) + so the global rescaling does not smear the border difference over the whole patch. + """ + volume = self.tiny_volume() + shared = {"p": 1.0, "n_layers": 4, "kernel_sizes": (3,), "alpha_range": (1.0, 1.0), "out_norm": "none"} + reflected = self.apply(RandomGINGPU(padding_mode="reflect", **shared), volume, 1) + zeroed = self.apply(RandomGINGPU(padding_mode="zeros", **shared), volume, 1) + + interior = (slice(None), slice(None), slice(4, -4), slice(4, -4), slice(4, -4)) + self.assertTrue(bool(torch.equal(reflected[interior], zeroed[interior])), "the padding reached further in than the cascade can") + self.assertGreater(float((reflected - zeroed).abs().max()), 1.0, "the two padding modes produced the same border") + + +class TestTheCallContract(GinTestCase): + def test_it_runs_with_an_empty_params_dict(self): + """`params={}` -- no segmentation, no `batch_prob` -- is how it is called directly. + + `RandomChooseXTransformsGPU` dispatches to `apply_transform` that way, and so does + every test here. + """ + volume = self.tiny_volume() + transform = RandomGINGPU(p=1.0) + self.assertIsImageLike(transform.apply_transform(volume.clone(), {}, transform.flags), volume, "RandomGINGPU") + + def test_it_does_not_write_into_the_callers_tensor(self): + """Pinned here as well as in the sweep: the gather and scatter is where this is lost.""" + volume = self.tiny_volume() + keep = volume.clone() + transform = RandomGINGPU(p=1.0) + transform.apply_transform(volume, {}, transform.flags) + self.assertTrue(bool(torch.equal(volume, keep)), "apply_transform modified its argument") + + def test_region_selection_keeps_the_complement_of_the_segmentation(self): + """`in_seg=1.0` confines the result to the labelled voxels and nothing else moves.""" + volume = self.tiny_volume() + seg = self.tiny_seg() + out = self.apply(RandomGINGPU(p=1.0, in_seg=1.0), volume, seg=seg) + outside = seg == 0 + self.assertTrue(bool(torch.equal(out[outside], volume[outside])), "the transform leaked outside the segmentation") + self.assertFalse(bool(torch.equal(out, volume)), "in_seg=1.0 left the whole patch alone") diff --git a/unit_tests/test_transforms_gpu.py b/unit_tests/test_transforms_gpu.py index d370f68..9318221 100644 --- a/unit_tests/test_transforms_gpu.py +++ b/unit_tests/test_transforms_gpu.py @@ -25,6 +25,7 @@ "smauglab.transforms.gpu.contrast", "smauglab.transforms.gpu.spatial", "smauglab.transforms.gpu.fromSeg", + "smauglab.transforms.gpu.gin", "smauglab.transforms.gpu.domain_transfer", "smauglab.transforms.gpu.palette.transform", ]