From af04f28dae36c38a0b6de5feff7c46feda197bf1 Mon Sep 17 00:00:00 2001 From: iback Date: Tue, 6 Oct 2026 11:33:03 +0000 Subject: [PATCH] feat: GIN, a cascade of random convolutions for domain generalisation GIN pushes the image through a shallow convolutional network whose weights are drawn afresh on every call, blends the result back towards the original at a per-sample alpha, and rescales so the output 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. 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 TMI 42(4), 1095-1106. DOI 10.1109/TMI.2022.3224067 Ported from the authors' own 3-D implementation (`models/imagefilter3d.py` in cheng-01037/Causality-Medical-Image-Domain-Generalization) and the cleaner nnU-Net-oriented restatement in `dg_tta/gin.py` (multimodallearning/DG-TTA). Both are MIT licensed. `RandomRandConvGPU` was the closest thing here and is single-layer RandConv: one random kernel, no non-linearity, no output renormalisation. GIN is its multi-layer non-linear generalisation, so it gets its own AugId -- sharing RAND_CONV's would have overwritten that transform's matrix cell -- and sits next to it in PIPELINE_ORDER. Two deliberate departures from upstream, both argued at their call sites: * Reflect padding rather than zeros. Every other convolution in this package reflect-pads, and a zero-padded 3x3x3 kernel darkens the patch border, which an nnU-Net patch has no business having: it is an interior crop, not a whole image. `padding_mode="zeros"` restores upstream exactly, and `test_zero_padding_changes_only_the_border` pins that the choice reaches exactly as far into the patch as the cascade can. * float32 as a floor for the cascade and the norm. The trainer runs the GPU pipeline inside `autocast`, so `F.conv3d` would hand back float16; the kernels are raw N(0, 1) with no fan-in scaling, and simulating the stack without the forcing, `||mixed||_F` passed float16's 65504 in 23 of 150 draws at the defaults on a 1x128^3 patch, 48 of 150 on 2x192^3, and 105 of 150 with kernel_sizes=(1, 3, 5, 7). The failure is not NaN and that is the whole problem: the norm comes back `inf`, `1 / (inf + 1e-5)` is exactly 0, and the output is an all-zero patch that `torch.isfinite(...).all()` reports as fine. Neither `_select_and_check` nor the trainer's non-finite-loss guard can see it, and the network trains on a blank volume. `promote_types` rather than `.float()` so a float64 input keeps its precision. The unscaled kernels are safe because of the Frobenius renormalisation, which is exact rather than statistical: the per-sample output RMS matches the input's to within 4.1e-07 relative for every combination of n_layers in {1, 2, 4, 8} and kernel_sizes in {(1,), (3,), (1,3), (1,3,5,7)}. That is why GIN needs no `1/sqrt(k**3)` where RandConv does. 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 about 1.3x across those settings. `test_without_the_frobenius_norm_the_scale_does_track_the_configuration` is the control -- turning the renormalisation off lets the scale run away by 19197x over the same sweep. The channels named by `apply_to_channel` are mapped jointly, by one network with that many inputs and outputs, because cross-channel mixing is the published method. For single-channel data, which is the usual nnU-Net case, that is identical to a per-channel loop. It does make a duplicate entry change the network's width rather than merely applying it twice, so `(0, 0)` is now rejected at construction. There is no `mix_prob`, unlike the sibling transforms: `alpha_range` already is the blend back towards the original, and a `mix_prob` on top would blend twice. Augmentation behaviour is unchanged. `transform_params_gpu.json` gains `"RandomGINGPU": {"p": 0.0}` so the default config names it and a sweep only has to flip one number; the config has no `pipeline` section so it runs SEQUENTIAL, and `_leaf_parameters` short-circuits p == 0 without touching the RNG. Verified by building the pipeline from the config before and after under one seed: bitwise identical, max abs diff 0.0, against 2.886 for the same entry at p=1.0. Only the content hash moves, 49138fe4 -> 1c7447e7. Co-Authored-By: Claude Opus 5 --- README.md | 1 + smauglab/configs/all_augmentations.json | 26 ++ smauglab/configs/transform_params_gpu.json | 3 + smauglab/registry.py | 8 + smauglab/transforms/__init__.py | 1 + smauglab/transforms/gpu/gin.py | 340 +++++++++++++++++++++ unit_tests/test_gin.py | 323 ++++++++++++++++++++ unit_tests/test_transforms_gpu.py | 1 + 8 files changed, 703 insertions(+) create mode 100644 smauglab/transforms/gpu/gin.py create mode 100644 unit_tests/test_gin.py diff --git a/README.md b/README.md index cb2b381..e80f746 100644 --- a/README.md +++ b/README.md @@ -232,6 +232,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", ]