diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index f5a39fa..5c2baef 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -46,6 +46,7 @@ data on disk, so it is fast enough to gate every pull request. | `test_configs.py` | Every shipped config parses, builds a pipeline, runs a forward pass, and is reproducible under a fixed seed | | `test_transforms_gpu.py` | Each GPU transform in isolation | | `test_packaging.py` | Builds the real wheel and checks its contents | +| `test_consistency_loss.py` | The `loss.consistency` term, the geometric/intensity pipeline split it needs, and its guardrails | Tests are `unittest.TestCase` subclasses, so they run under either runner: diff --git a/README.md b/README.md index 4755ed2..74e83ff 100644 --- a/README.md +++ b/README.md @@ -187,6 +187,61 @@ Scripts developped in this repository use JSON files to specify image and segmen To track parameters used during data augmentation, JSON files are also used: see this [example](https://github.com/neuropoly/SmaugLab/blob/16653a84e031c40e25a72e946c2724494606b21c/smauglab/configs/transform_params.json) +### Domain-generalization consistency loss + +An optional second loss term for the nnU-Net trainer. The augmented patch and the +*unaugmented* one both go through the network, and their predictions are compared with +a soft Dice: the two images show the same anatomy, so they should get the same +segmentation, whatever the appearance augmentation did to the intensities. + +The reference is taken **after the geometric transforms and before the intensity ones**. +Geometry (`RandomFlipTransformGPU`, `RandomAffineGPU`) runs once and is shared by both +branches, so the two images are aligned voxel for voxel and one segmentation is correct +for both. The term is off unless a config asks for it: + +```json +{ + "loss": { + "consistency": { + "enabled": true, + "weight": 0.5, + "detach_reference": true + } + }, + "GPU": { "...": {} } +} +``` + +| key | default | what it does | +| --- | --- | --- | +| `enabled` | `false` | Turns the term on. Needs a `GPU` section; a CPU-only config is rejected, because there the dataloader has already applied everything and there is no unaugmented image left. | +| `weight` | `0.5` | Multiplier on the term. `0` measures and logs it without changing training, which is the ablation worth running first. | +| `detach_reference` | `true` | Runs the reference forward under `no_grad`. With `false` both branches are in the graph: roughly double the activation memory, and the term becomes minimisable by making the *clean* prediction worse. | + +Costs one extra forward pass per step. Only the full-resolution deep-supervision head +takes part; the lower-resolution heads keep nnU-Net's own weighting for the supervised +term alone. + +Two things to know when reading the numbers: + +* `train_loss` now carries both terms, so it is not comparable with a run that had the + term off. The epoch log prints the component on its own line + (`consistency: soft dice ... contributing ... to train_loss`). +* The term is a Dice against a *soft* target, so it does not reach `-1` when the two + branches agree; its minimiser is the reference's hard argmax, which means it sharpens + predictions as well as aligning them. + +`RandomCropTransformGPU` interacts with it. The crop occludes part of the field of view +rather than moving voxels, so it is grouped `OCC` rather than `GEO` and stays behind the +intensity stage -- in front of it, the intensity transforms would paint texture into the +region it discarded. It is therefore not applied to the reference branch, and in the +batches where it fires the consistency term asks the network to agree with the reference +inside the occluded region, where the supervised target says background. Nothing masks +this; the trainer prints a warning at startup naming the probability. Set its `p` to `0` +to avoid it, which is what every shipped config other than `transform_params_gpu.json`, +`transform_params_gpu_default01-23_CropTransform.json` and +`transform_params_gpu_palette-em.json` already does. + ### Paper configs Two shipped configs are the setups behind the paper's tables. `smauglab hash ` @@ -220,7 +275,7 @@ column -- see "Run Monai training with SmaugLab augmentations" above. Regenerate | --- | --- | --- | --- | | flip | GEO | `RandomFlipTransformGPU` | — | | affine | GEO | `RandomAffineGPU` | — | -| crop | GEO | `RandomCropTransformGPU` | — | +| crop | OCC | `RandomCropTransformGPU` | — | | spatial | GEO | — | `SpatialTransform` | | gaussian_noise | GE | `RandomGaussianNoiseGPU` | `GaussianNoiseTransform` | | gaussian_blur | GE | `RandomGaussianBlurGPU` | `GaussianBlurTransform` | diff --git a/smauglab/cli.py b/smauglab/cli.py index 0cf470d..35a9454 100644 --- a/smauglab/cli.py +++ b/smauglab/cli.py @@ -232,7 +232,7 @@ def build_parser() -> argparse.ArgumentParser: p = sub.add_parser("list", help="registered augmentations, in pipeline order") p.add_argument("--backend", choices=["gpu", "cpu"]) - p.add_argument("--group", choices=["geo", "ge", "ta"]) + p.add_argument("--group", choices=["geo", "ge", "ta", "occ"]) p.add_argument("--json", action="store_true") p.set_defaults(func=cmd_list) diff --git a/smauglab/config.py b/smauglab/config.py index 93bd7ca..494042e 100644 --- a/smauglab/config.py +++ b/smauglab/config.py @@ -6,7 +6,8 @@ "_comment": "...", // any key starting with '_' is ignored "GPU": { "": { }, ... }, "CPU": { ... }, - "pipeline": { "random_choose": { ... } } + "pipeline": { "random_choose": { ... } }, + "loss": { "consistency": { "enabled": true, "weight": 0.5 } } } A key is a class name, exactly, and a parameter is a constructor argument, exactly. Both are @@ -25,8 +26,10 @@ import difflib import functools import json +from collections.abc import Mapping from enum import Enum from pathlib import Path +from types import MappingProxyType from typing import Any from smauglab import registry @@ -93,11 +96,35 @@ def validate_section(section: dict, backend: Backend, *, source: str = " #: Top-level keys that are not augmentation sections. -RESERVED_SECTIONS = ("pipeline",) +RESERVED_SECTIONS = ("pipeline", "loss") #: Keys the `pipeline` section may hold. PIPELINE_KEYS = ("mode", "order", "random_choose", "same_on_batch") +#: Loss terms the `loss` section can configure. The supervised term is nnU-Net's own +#: and is not reachable from here; these are the ones SmaugLab adds on top of it. +LOSS_KEYS = ("consistency",) + +#: Every setting each loss term accepts, with its default. Two jobs in one table: it is +#: the accepted-key list validation checks against, and it is the defaults +#: `SmaugConfig.loss_options` merges in, so a default cannot drift between the trainer, +#: the startup summary and the tests. +LOSS_DEFAULTS: Mapping[str, Mapping[str, Any]] = MappingProxyType( + { + "consistency": MappingProxyType( + { + # Off, so every config predating this section behaves exactly as before. + "enabled": False, + "weight": 0.5, + # False keeps the reference forward in the graph: roughly double the + # activation memory, and the term becomes minimisable by degrading the + # reference prediction instead of improving the augmented one. + "detach_reference": True, + } + ) + } +) + #: What to do about a config that has no backend section. Quoted by both errors that #: can report it, so the two cannot drift apart. SECTION_HINT = "move the augmentation blocks into a 'GPU' or 'CPU' section, and spell each key as its class name." @@ -137,6 +164,20 @@ def pipeline_options(self, name: str) -> dict[str, Any]: """ return copy.deepcopy(self.payload.get("pipeline", {}).get(name, {})) + def loss_options(self, name: str) -> dict[str, Any]: + """One loss term's settings, defaults filled in. Always a dict, always a copy. + + The defaults from `LOSS_DEFAULTS` are merged here rather than left to each + caller's `.get(key, default)`: the trainer, the startup summary and the tests + all have to get the same answer out of a partially specified block. + """ + block = self.payload.get("loss", {}).get(name, {}) + return {**LOSS_DEFAULTS.get(name, {}), **copy.deepcopy(block)} + + def consistency_enabled(self) -> bool: + """Whether the augmented/reference consistency term is on.""" + return bool(self.loss_options("consistency")["enabled"]) + def pipeline_mode(self) -> PipelineMode: """How the GPU pipeline arranges its transforms. @@ -198,6 +239,7 @@ def validate(self) -> None: problems.append(f"no GPU or CPU section; this looks like a pre-registry config. {SECTION_HINT}") problems.extend(self._pipeline_problems()) + problems.extend(self._loss_problems()) for backend in Backend: section = self.payload.get(backend.value) @@ -248,6 +290,71 @@ def _pipeline_problems(self) -> list[str]: problems.append(f"pipeline.mode: unknown mode {mode!r}. Accepted: {', '.join(valid)}.{hint}") return problems + def _loss_problems(self) -> list[str]: + """Check the `loss` section, the same way `_pipeline_problems` checks its own. + + A typo'd setting has to be an error rather than an ignored key: a run whose + config says `"detatch_reference": false` would otherwise train with the + reference detached and record a config claiming the opposite, and that config + is the only provenance a finished run carries. + """ + loss = self.payload.get("loss") + if loss is None: + return [] + if not isinstance(loss, dict): + return ["loss: expected an object"] + + problems = [] + for key in loss: + if key.startswith("_") or key in LOSS_KEYS: + continue + close = difflib.get_close_matches(key, LOSS_KEYS, n=2, cutoff=0.6) + hint = f" Did you mean: {', '.join(close)}?" if close else "" + problems.append(f"loss: unknown term {key!r}. Accepted: {', '.join(LOSS_KEYS)}.{hint}") + + for term in LOSS_KEYS: + block = loss.get(term) + if block is None: + continue + if not isinstance(block, dict): + problems.append(f"loss.{term}: expected an object of settings") + continue + accepted = tuple(LOSS_DEFAULTS[term]) + for key in block: + if key.startswith("_") or key in accepted: + continue + close = difflib.get_close_matches(key, accepted, n=2, cutoff=0.6) + hint = f" Did you mean: {', '.join(close)}?" if close else "" + problems.append(f"loss.{term}: unknown setting {key!r}. Accepted: {', '.join(accepted)}.{hint}") + problems.extend(_consistency_setting_problems(term, block)) + return problems + + +def _consistency_setting_problems(term: str, block: dict) -> list[str]: + """Type-check one loss term's settings. Split out so the checks read as a list. + + Presence is tested with `in`, not against None. A JSON `null` is not a request for + the default: `loss_options` merges the block over `LOSS_DEFAULTS`, so a null would + *replace* the default with None and the trainer's `float(None)` would raise -- or, + worse for `enabled`, `bool(None)` would quietly turn the term off. + """ + problems = [ + f"loss.{term}.{key}: expected true or false, got {block[key]!r}" + for key in ("enabled", "detach_reference") + if key in block and not isinstance(block[key], bool) + ] + + if "weight" in block: + weight = block["weight"] + # `isinstance(True, int)` is True, so bool has to be excluded explicitly. + if isinstance(weight, bool) or not isinstance(weight, (int, float)): + problems.append(f"loss.{term}.weight: expected a number, got {weight!r}") + elif weight < 0: + # 0 is allowed: it is the "pay the second forward, measure the term, change + # nothing" ablation, and the epoch summary still reports it. + problems.append(f"loss.{term}.weight: expected a non-negative number, got {weight!r}") + return problems + @functools.lru_cache(maxsize=8) def load_config(path: str) -> SmaugConfig: diff --git a/smauglab/registry.py b/smauglab/registry.py index 36805bc..2ad41f2 100644 --- a/smauglab/registry.py +++ b/smauglab/registry.py @@ -65,15 +65,27 @@ class Backend(str, Enum): class AugType(str, Enum): """Augmentation strength class. - Member names are verbatim from segtransferaug's `AUG2GROUP`, which this + GEO, GE and TA are verbatim from segtransferaug's `AUG2GROUP`, which this replaces, so downstream `AugType.TA` keeps resolving. The GPU random-order pipelines bucket transforms by this to decide what goes inside each `RandomChooseXTransformsGPU`. + + OCC has no `AUG2GROUP` ancestor and is GEO for every purpose the pipeline + has: `build.SEQUENTIAL_GROUPS` runs it in the same leading sequential run and + `PIPELINE_ORDER` is untouched. It exists so the consistency loss can decline + to hoist it. That loss runs the GEO transforms once and shares the result with + its reference branch (`build.SHARED_GROUPS`), and a transform that *blanks* + part of the volume rather than moving it must not go in that stage: the + intensity transforms that follow would paint texture into the region it + discarded. The split rests on GEO being a prefix of the pipeline, which is + also why the two cannot be one group -- `RandomCropTransformGPU` sits near + the end of `PIPELINE_ORDER`. """ - GEO = "GEO" # geometric: applied in sequence, never bucketed + GEO = "GEO" # geometric: moves voxels, applied in sequence, never bucketed GE = "GE" # general enhancement, usually light TA = "TA" # transfer augmentation, usually heavy + OCC = "OCC" # occlusion: sequential like GEO, but discards field of view class AugId(str, Enum): diff --git a/smauglab/trainers/nnUNetTrainerDAExt.py b/smauglab/trainers/nnUNetTrainerDAExt.py index a9dac85..f7af87e 100644 --- a/smauglab/trainers/nnUNetTrainerDAExt.py +++ b/smauglab/trainers/nnUNetTrainerDAExt.py @@ -32,15 +32,16 @@ from batchgeneratorsv2.transforms.utils.deep_supervision_downsampling import DownsampleSegForDSTransform from batchgeneratorsv2.transforms.utils.remove_label import RemoveLabelTansform from batchgeneratorsv2.transforms.utils.seg_to_regions import ConvertSegmentationToRegionsTransform +from nnunetv2.training.loss.dice import MemoryEfficientSoftDiceLoss from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer -from nnunetv2.utilities.helpers import dummy_context +from nnunetv2.utilities.helpers import dummy_context, softmax_helper_dim1 from torch import autocast from smauglab import configs from smauglab.config import load_config -from smauglab.registry import Backend +from smauglab.registry import AugType, Backend, InvalidConfigError, entries from smauglab.trainers.utils import DownsampleSegForDSTransformCustom, nnunet_tail_transforms -from smauglab.transforms.build import build_cpu_pipeline, build_gpu_pipeline +from smauglab.transforms.build import SHARED_GROUPS, build_cpu_pipeline, build_gpu_stages from smauglab.transforms.gpu.base import AugmentationSequentialCustom #: Env var naming the config to train with. @@ -73,6 +74,10 @@ #: `torch.amp.GradScaler`'s default `init_scale`. Only used to warn about a floor at or #: above it, which would undo every backoff the scaler ever makes. GRAD_SCALER_INIT_SCALE = 2.0**16 +#: nnU-Net's own dice smoothing, from `nnUNetTrainer._build_loss`. Repeated here so the +#: supervised dice and the consistency dice that end up in the same sum are computed +#: identically; a different epsilon would make the two terms incomparable. +DICE_SMOOTH = 1e-5 def resolve_config_path() -> str: @@ -170,6 +175,63 @@ def _has_gpu_augmentations(config) -> bool: return bool(config.names(Backend.GPU)) +def _active(section: dict, names: list[str]) -> list[str]: + """Of `names`, the ones the section enables with a non-zero probability.""" + return [name for name in names if isinstance(section.get(name), dict) and section[name].get("p", 1.0) != 0] + + +def _consistency_problems(config) -> list[str]: + """Why this config cannot train with the consistency term. Empty means it can.""" + if not _has_gpu_augmentations(config): + return [ + "loss.consistency needs the intensity augmentations in the GPU section: the reference " + "image is the GPU pipeline's geometric stage, and with a CPU-only config the dataloader " + "has already applied everything, so there is nothing left to compare against. Move the " + "augmentations into a 'GPU' section, or set loss.consistency.enabled to false." + ] + return [] + + +def _consistency_warnings(config) -> list[str]: + """Things about this config that weaken the term without making it wrong.""" + warnings_out = [] + gpu = config.section(Backend.GPU) + cpu = config.section(Backend.CPU) + + shared = [entry.name for entry in entries(backend=Backend.GPU) if entry.group in SHARED_GROUPS] + if not _active(gpu, shared): + warnings_out.append( + f"consistency: no geometric augmentation is active in the GPU section " + f"({', '.join(shared)}), so the shared stage is empty and the reference is the " + f"dataloader's patch verbatim. The term still works -- nothing but the geometric " + f"group touches the mask -- but this is usually a config mistake." + ) + + # CPU GEO is fine and must not be warned about: it runs before the batch reaches + # train_step, so both branches inherit it identically. Only intensity work on the + # CPU side makes the "unaugmented" reference a misnomer. + intensity = [entry.name for entry in entries(backend=Backend.CPU) if entry.group in (AugType.GE, AugType.TA)] + offenders = _active(cpu, intensity) + if offenders: + warnings_out.append( + f"consistency: the CPU section applies intensity augmentations in the dataloader " + f"({', '.join(offenders)}), so the reference is already augmented and the term compares " + f"two differently augmented views rather than augmented against clean. Still a valid " + f"objective, just not the one the name implies." + ) + + occluders = [entry.name for entry in entries(backend=Backend.GPU) if entry.group is AugType.OCC] + for name in _active(gpu, occluders): + probability = gpu[name].get("p", 1.0) + warnings_out.append( + f"consistency: {name} fires at p={probability} and is not applied to the reference " + f"branch, so in those batches the term asks the network to agree with the reference " + f"inside the occluded region, where the supervised target says background. The two " + f"terms disagree there and nothing masks it; set its p to 0 to avoid it." + ) + return warnings_out + + class nnUNetTrainerDAExtGPU(nnUNetTrainer): """nnU-Net trainer driven entirely by a SmaugLab config. @@ -183,28 +245,83 @@ def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dic json_path = resolve_config_path() config = load_config(json_path) + consistency = config.loss_options("consistency") + split = bool(consistency["enabled"]) + if split: + problems = _consistency_problems(config) + if problems: + raise InvalidConfigError(config.source, problems) + # Built only when the config actually asks for GPU augmentations, so a # CPU-only config costs nothing per step and train_step stays a no-op. + # + # `self.transforms` is the augmented branch either way, so train_step's existing + # path is untouched when the consistency term is off. `self.transforms_shared` + # holds the geometric stage only when it is on; both containers get the mask, + # because the segmentation-conditioned transfer augmentations read it through the + # SegmentationRef that AugmentationSequentialCustom injects per call. self.transforms: AugmentationSequentialCustom | None = None + self.transforms_shared: AugmentationSequentialCustom | None = None if _has_gpu_augmentations(config): - self.transforms = AugmentationSequentialCustom( - *build_gpu_pipeline( - config.section(Backend.GPU), - mode=config.pipeline_mode(), - options=config.pipeline_options("random_choose"), - source=config.source, - order_source=config.order_source(), - ), - data_keys=["input", "mask"], - # See AugTransformsGPU: True forces the flag onto every child and - # spends each transform's `p` once for the whole batch. - same_on_batch=True if config.same_on_batch() else None, - ).to(self.device) + shared, augmented = build_gpu_stages( + config.section(Backend.GPU), + mode=config.pipeline_mode(), + options=config.pipeline_options("random_choose"), + source=config.source, + order_source=config.order_source(), + split_shared=split, + ) + + def container(*stage) -> AugmentationSequentialCustom: + return AugmentationSequentialCustom( + *stage, + data_keys=["input", "mask"], + # See AugTransformsGPU: True forces the flag onto every child and + # spends each transform's `p` once for the whole batch. + same_on_batch=True if config.same_on_batch() else None, + ).to(self.device) + + if shared: + self.transforms_shared = container(*shared) + self.transforms = container(*augmented) + + # The criterion is built here rather than in `_build_loss`: it is a second term + # beside `self.loss`, not a replacement for it, and everything it needs + # (`label_manager`, `configuration_manager`, `is_ddp`) is already assigned by + # nnUNetTrainer.__init__. + self.consistency: MemoryEfficientSoftDiceLoss | None = None + self.consistency_weight = 0.0 + self.consistency_detach = True + self.consistency_nonlin = softmax_helper_dim1 + if split: + self.consistency_weight = float(consistency["weight"]) + self.consistency_detach = bool(consistency["detach_reference"]) + # Regions mean a sigmoid head and a one-hot target, as in DC_and_BCE_loss; + # otherwise softmax and do_bg=False. Mirrors `_build_loss` so the two dice + # terms in `l` disagree about nothing. + self.consistency_nonlin = torch.sigmoid if self.label_manager.has_regions else softmax_helper_dim1 + self.consistency = MemoryEfficientSoftDiceLoss( + apply_nonlin=self.consistency_nonlin, + batch_dice=self.configuration_manager.batch_dice, + do_bg=self.label_manager.has_regions, + smooth=DICE_SMOOTH, + ddp=self.is_ddp, + ) print(f"Using SmaugLab transforms from: {json_path}") print( f" CPU: {len(config.names(Backend.CPU))} augmentations, GPU: {len(config.names(Backend.GPU))}, mode: {config.pipeline_mode().value}" ) + if split: + print( + f" loss.consistency: on (weight {self.consistency_weight}, " + f"detach_reference {self.consistency_detach}), shared geometric stage: " + f"{len(self.transforms_shared) if self.transforms_shared is not None else 0} transforms" + ) + for warning in _consistency_warnings(config): + print(f" {warning}") + else: + print(" loss.consistency: off") _record_config(json_path, os.path.join(self.output_folder, "transform_params_used_for_training.json")) @@ -227,6 +344,12 @@ def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dic ) print(f"AMP gradient-scale floor set to {self._min_grad_scale} via {MIN_GRAD_SCALE_ENV}{note}") + # Accumulated as a device tensor, not a float: reading the term back per step + # would add a host sync train_step does not otherwise pay until its last line. + self._consistency_sum: torch.Tensor | None = None + self._consistency_steps = 0 + self._consistency_last: torch.Tensor | None = None + @staticmethod def get_training_transforms( patch_size: Union[np.ndarray, tuple[int, ...]], @@ -339,6 +462,16 @@ def train_step(self, batch: dict) -> dict: # GPU augmentations run on the full-resolution data/target, then the # deep-supervision targets are built from the *augmented* mask. A CPU-only # config builds no GPU pipeline and the dataloader already produced them. + reference = None + if self.transforms_shared is not None: + # The geometric stage, run once. Both branches see this output, so the + # one segmentation that follows is correct for both of them. + data, target = self.transforms_shared(data, target) + # Not a clone. unit_tests/test_no_inplace_mutation.py pins that no + # registered transform writes into the tensor it was handed, so the + # second stage cannot reach back into this one, and a clone here would + # cost a whole extra volume on every step. + reference = data if self.transforms is not None: data, target = self.transforms(data, target) @@ -350,6 +483,8 @@ def train_step(self, batch: dict) -> dict: output = self.network(data) # del data l = self.loss(output, target) + if reference is not None and self.consistency is not None: + l = l + self.consistency_weight * self._consistency_term(output, reference) if self.grad_scaler is not None: self.grad_scaler.scale(l).backward() @@ -376,6 +511,10 @@ def train_step(self, batch: dict) -> dict: loss_cpu = l.detach().cpu() if not bool(torch.isfinite(loss_cpu)): self._report_nonfinite_loss(loss_cpu, data, target, output, batch.get("keys")) + # nnU-Net's `on_train_epoch_end` reads only "loss", so the extra key is inert + # upstream and the override below averages it. Absent when the term is off. + if self._consistency_last is not None: + return {"loss": loss_cpu.numpy(), "consistency": float(self._consistency_last)} return {"loss": loss_cpu.numpy()} def _guard_grad_scale(self) -> None: @@ -417,6 +556,45 @@ def _guard_grad_scale(self) -> None: elif self._scale_floor_hits_this_epoch == NAN_REPORTS_PER_EPOCH + 1: self.print_to_log_file("AMP gradient scale: further resets suppressed this epoch; see the epoch summary") + def _consistency_term(self, output, reference: torch.Tensor) -> torch.Tensor: + """Soft dice between the augmented and the reference prediction. + + Only the full-resolution head takes part. The supervised term spreads over the + deep-supervision heads with the weights `nnUNetTrainer._build_loss` builds -- + including its `weights[-1] = 0` and the DDP-dependent `1e-6` variant of it -- + and re-deriving them here would copy a decision that belongs upstream. Head 0 + already carries most of the supervised weight, and it is the resolution the + shared geometric stage is exact at. + + The reference is handed over as a *probability map*, not a label map: + `MemoryEfficientSoftDiceLoss` takes its `x.shape == y.shape` branch and uses it + as-is. Its sibling `SoftDiceLoss` cannot -- `get_tp_fp_fn_tn` does `~y_onehot`, + which is a TypeError on a float tensor -- which is why this is the one of the + two nnU-Net dice classes used here. + + Note that with a soft target the term does not bottom out at -1 when the two + branches agree: its minimiser over the augmented prediction is the *hard* argmax + of the reference, so it aligns the two branches and sharpens them at once. The + logged number is therefore a dice against a soft target, not an agreement score. + """ + criterion = self.consistency + if criterion is None: # pragma: no cover -- train_step guards the call + raise RuntimeError("_consistency_term needs a criterion; loss.consistency is not enabled") + + with torch.no_grad() if self.consistency_detach else dummy_context(): + head = self.network(reference) + head = head[0] if isinstance(head, (list, tuple)) else head + probabilities = self.consistency_nonlin(head) + + prediction = output[0] if isinstance(output, (list, tuple)) else output + term = criterion(prediction, probabilities) + + detached = term.detach() + self._consistency_last = detached + self._consistency_sum = detached if self._consistency_sum is None else self._consistency_sum + detached + self._consistency_steps += 1 + return term + def _report_nonfinite_loss(self, loss, data, target, output, keys) -> None: """Say *why* the loss went non-finite, in the training log. @@ -449,18 +627,26 @@ def _report_nonfinite_loss(self, loss, data, target, output, keys) -> None: elif not bool(torch.isfinite(head).all()): dead = any(not bool(torch.isfinite(p).all()) for p in self.network.parameters()) verdict, detail = "network", f"output_nonfinite=True weights_nonfinite={dead}" + elif self._consistency_last is not None and not bool(torch.isfinite(self._consistency_last)): + # Without this branch a NaN out of the reference forward is reported as + # verdict=loss, which points the reader at nnU-Net's criterion -- the wrong file. + verdict, detail = "consistency", "data and output are finite; the consistency term is not" else: verdict, detail = "loss", "data and output are finite" + consistency = f" consistency={float(self._consistency_last)}" if self._consistency_last is not None else "" self.print_to_log_file( f"non-finite loss: value={float(loss)} verdict={verdict} epoch={self.current_epoch} " - f"{detail} target_fg_voxels={int((seg > 0).sum())} keys={list(keys) if keys is not None else 'n/a'}" + f"{detail}{consistency} target_fg_voxels={int((seg > 0).sum())} " + f"keys={list(keys) if keys is not None else 'n/a'}" ) def on_train_epoch_start(self): super().on_train_epoch_start() self._nan_steps_this_epoch = 0 self._scale_floor_hits_this_epoch = 0 + self._consistency_sum = None + self._consistency_steps = 0 def on_train_epoch_end(self, train_outputs: list[dict]): super().on_train_epoch_end(train_outputs) @@ -480,3 +666,14 @@ def on_train_epoch_end(self, train_outputs: list[dict]): f"({self._scale_floor_hits_total} since the run started); " f"without the floor the run would be unrecoverable by now" ) + if self._consistency_steps and self._consistency_sum is not None: + # `nnUNetLogger.log` asserts the key already exists in its + # `my_fantastic_logging` dict, so a new curve cannot be added without + # patching upstream. Printed beside the epoch's training section instead, + # like the non-finite-loss summary above -- and it is the only way to read + # the epoch line, whose train_loss now carries both terms. + mean = float(self._consistency_sum) / self._consistency_steps + self.print_to_log_file( + f"consistency: soft dice {-mean:.4f} over {self._consistency_steps} batches, " + f"weight {self.consistency_weight}, contributing {self.consistency_weight * mean:+.4f} to train_loss" + ) diff --git a/smauglab/transforms/build.py b/smauglab/transforms/build.py index 63fe621..664e05e 100644 --- a/smauglab/transforms/build.py +++ b/smauglab/transforms/build.py @@ -14,6 +14,7 @@ from __future__ import annotations +import itertools from typing import Any from smauglab import registry @@ -76,6 +77,16 @@ def build_transforms( return sorted(built, key=lambda pair: registry.pipeline_position(pair[1])) +#: Groups that run in sequence ahead of any shuffled bucket. OCC is listed beside GEO +#: deliberately: the crop has always run here, and keeping it here is what makes every +#: built pipeline identical to the one before `AugType.OCC` existed. +SEQUENTIAL_GROUPS: tuple[AugType, ...] = (AugType.GEO, AugType.OCC) + +#: Groups `build_gpu_stages` may hoist into the stage both branches of the consistency +#: loss share. OCC is deliberately absent -- see `AugType`. +SHARED_GROUPS: tuple[AugType, ...] = (AugType.GEO,) + + def _buckets(built: list[tuple[Any, AugEntry]], *, hoist_sequential: bool) -> tuple[list, list, list]: """Split into (sequential, transfer, general-enhancement). @@ -88,7 +99,7 @@ def _buckets(built: list[tuple[Any, AugEntry]], *, hoist_sequential: bool) -> tu """ sequential, transfer, enhancement = [], [], [] for transform, entry in built: - if entry.group is AugType.GEO or (hoist_sequential and entry.force_sequential): + if entry.group in SEQUENTIAL_GROUPS or (hoist_sequential and entry.force_sequential): sequential.append(transform) elif entry.group is AugType.TA: transfer.append(transform) @@ -97,16 +108,33 @@ def _buckets(built: list[tuple[Any, AugEntry]], *, hoist_sequential: bool) -> tu return sequential, transfer, enhancement -def build_gpu_pipeline( - section: dict, - *, - mode: PipelineMode = PipelineMode.SEQUENTIAL, - options: dict[str, Any] | None = None, - source: str = "", - order_source: OrderSource = OrderSource.REGISTRY, -) -> list[Any]: - """The GPU transform list for a given pipeline mode.""" - built = build_transforms(section, Backend.GPU, source=source, order_source=order_source) +def _split_shared(built: list[tuple[Any, AugEntry]], *, source: str) -> tuple[list, list]: + """Partition into (shared by both branches, augmented branch only). + + The shared stage has to be a *prefix*. Hoisting a transform past one that stays + behind silently reorders the pipeline for every config naming both, and the + consistency loss would then be comparing a reference built by a pipeline nobody + configured. Today GPU `SHARED_GROUPS` is {flip, affine} at `PIPELINE_ORDER` + positions 0 and 1, so the prefix is free -- but a later registration, or + `pipeline.order: config`, can break it, and this is the only place that notices. + """ + shared = [pair for pair in built if pair[1].group in SHARED_GROUPS] + prefix = list(itertools.takewhile(lambda pair: pair[1].group in SHARED_GROUPS, built)) + if len(prefix) != len(shared): + late = [entry.name for _, entry in shared[len(prefix) :]] + raise InvalidConfigError( + source, + [ + f"the consistency loss needs the shared stage to be a prefix of the pipeline, but " + f"{', '.join(late)} run(s) after a transform that stays in the augmented branch. " + f"Reorder the pipeline, or disable loss.consistency." + ], + ) + return prefix, built[len(prefix) :] + + +def _assemble(built: list[tuple[Any, AugEntry]], *, mode: PipelineMode, options: dict[str, Any] | None) -> list[Any]: + """Arrange already-instantiated transforms according to the pipeline mode.""" if mode is PipelineMode.SEQUENTIAL: return [transform for transform, _ in built] @@ -142,6 +170,45 @@ def build_gpu_pipeline( return out +def build_gpu_stages( + section: dict, + *, + mode: PipelineMode = PipelineMode.SEQUENTIAL, + options: dict[str, Any] | None = None, + source: str = "", + order_source: OrderSource = OrderSource.REGISTRY, + split_shared: bool = False, +) -> tuple[list[Any], list[Any]]: + """The GPU pipeline, optionally split into (shared, augmented-only). + + With `split_shared=False` -- the default, and what `build_gpu_pipeline` asks for -- + the first list is empty and the second is the whole pipeline. With it True the + leading `SHARED_GROUPS` transforms move into the first list, which the consistency + loss runs once and hands to both branches so their segmentation references align. + + The two lists always concatenate to exactly what the unsplit call returns, in every + `PipelineMode`: the shared transforms are a prefix (`_split_shared` enforces it) and + `_assemble` only ever emits sequential groups at the front anyway. + """ + built = build_transforms(section, Backend.GPU, source=source, order_source=order_source) + shared, rest = _split_shared(built, source=source) if split_shared else ([], built) + return [transform for transform, _ in shared], _assemble(rest, mode=mode, options=options) + + +def build_gpu_pipeline( + section: dict, + *, + mode: PipelineMode = PipelineMode.SEQUENTIAL, + options: dict[str, Any] | None = None, + source: str = "", + order_source: OrderSource = OrderSource.REGISTRY, +) -> list[Any]: + """The GPU transform list for a given pipeline mode.""" + # `shared` is empty: build_gpu_stages was not asked to split. + shared, augmented = build_gpu_stages(section, mode=mode, options=options, source=source, order_source=order_source) + return shared + augmented + + def build_cpu_pipeline( section: dict, *, diff --git a/smauglab/transforms/gpu/spatial.py b/smauglab/transforms/gpu/spatial.py index 7b8b51b..a33e30b 100644 --- a/smauglab/transforms/gpu/spatial.py +++ b/smauglab/transforms/gpu/spatial.py @@ -614,7 +614,12 @@ def forward(self, batch_shape: tuple[int, ...], same_on_batch: bool = False) -> @register( aug_id=AugId.CROP, backend=Backend.GPU, - group=AugType.GEO, + # OCC, not GEO, although it runs in the same leading sequential position. GEO is + # what the consistency loss hoists into the stage its two branches share, and this + # one must not be hoisted: run before the intensity transforms, the flat slab it + # leaves behind gets texture painted into it -- the failure the `pad_value` note + # below describes -- which makes the occlusion no occlusion at all. See AugType. + group=AugType.OCC, ) class RandomCropTransformGPU(RigidAffineAugmentationBase3D): """Restrict the field of view to a random sub-box, filling the rest with `pad_value`. diff --git a/unit_tests/test_builder.py b/unit_tests/test_builder.py index f491e54..b487493 100644 --- a/unit_tests/test_builder.py +++ b/unit_tests/test_builder.py @@ -2,12 +2,43 @@ from __future__ import annotations -from typing import ClassVar +import hashlib +import json +from typing import Any, ClassVar -from smauglab.config import OrderSource, SmaugConfig -from smauglab.registry import Backend, InvalidConfigError -from smauglab.transforms.build import build_gpu_pipeline, build_transforms -from unit_tests.helpers import SmaugLabTestCase +from smauglab import registry +from smauglab.config import OrderSource, PipelineMode, SmaugConfig, load_config +from smauglab.registry import AugType, Backend, InvalidConfigError +from smauglab.transforms.build import SHARED_GROUPS, build_gpu_pipeline, build_gpu_stages, build_transforms +from unit_tests.helpers import SmaugLabTestCase, all_config_paths, gpu_config_paths, requires_external_asset + + +def _signature(transform: Any) -> Any: + """A structural fingerprint: class name, and a bucket's full contents.""" + inner = getattr(transform, "transforms_list", None) + if inner is None: + return type(transform).__name__ + return { + "bucket": type(transform).__name__, + "label": getattr(transform, "label", None), + "num": getattr(transform, "num_transforms", None), + "p": float(getattr(transform, "p", -1)), + "random_order": bool(getattr(transform, "random_order", False)), + "children": [_signature(child) for child in inner], + } + + +def _buildable_configs(): + """Every shipped config with a GPU section that builds on this machine.""" + for path in all_config_paths(): + if requires_external_asset(path): + continue + try: + config = load_config(str(path)) + except InvalidConfigError: + continue + if config.names(Backend.GPU): + yield path, config class TestOrderSource(SmaugLabTestCase): @@ -51,3 +82,145 @@ def test_every_problem_is_reported_at_once(self): with self.assertRaises(InvalidConfigError) as caught: build_transforms({"Nope": {}, "RandomFlipTransformGPU": {"bad": 1}}, Backend.GPU) self.assertEqual(len(caught.exception.problems), 2) + + +class TestSharedStageSplit(SmaugLabTestCase): + """`build_gpu_stages` splits the pipeline for the consistency loss without moving anything.""" + + def test_the_shared_stage_is_only_the_geometric_group(self): + geometric = set(registry.names(Backend.GPU, group=AugType.GEO)) + for path, config in _buildable_configs(): + with self.subTest(config=path.name): + shared, _ = build_gpu_stages( + config.section(Backend.GPU), + mode=config.pipeline_mode(), + options=config.pipeline_options("random_choose"), + order_source=config.order_source(), + split_shared=True, + ) + names = {type(t).__name__ for t in shared} + self.assertEqual(names, geometric & set(config.names(Backend.GPU))) + + def test_the_crop_never_enters_the_shared_stage(self): + """The whole point of AugType.OCC: the occlusion stays behind the intensity stage.""" + occluders = set(registry.names(Backend.GPU, group=AugType.OCC)) + self.assertIn("RandomCropTransformGPU", occluders) + for path, config in _buildable_configs(): + if not occluders & set(config.names(Backend.GPU)): + continue + for mode in PipelineMode: + with self.subTest(config=path.name, mode=mode.value): + shared, augmented = build_gpu_stages( + config.section(Backend.GPU), + mode=mode, + options=config.pipeline_options("random_choose"), + order_source=config.order_source(), + split_shared=True, + ) + self.assertFalse(occluders & {type(t).__name__ for t in shared}) + flat = json.dumps([_signature(t) for t in augmented]) + self.assertIn("RandomCropTransformGPU", flat) + + def test_the_two_stages_concatenate_to_the_unsplit_pipeline(self): + """The split is a partition, not a rearrangement, in every pipeline mode.""" + for path, config in _buildable_configs(): + for mode in PipelineMode: + with self.subTest(config=path.name, mode=mode.value): + kwargs = { + "mode": mode, + "options": config.pipeline_options("random_choose"), + "order_source": config.order_source(), + } + shared, augmented = build_gpu_stages(config.section(Backend.GPU), split_shared=True, **kwargs) + unsplit = build_gpu_pipeline(config.section(Backend.GPU), **kwargs) + self.assertEqual( + [_signature(t) for t in shared + augmented], + [_signature(t) for t in unsplit], + ) + + def test_not_splitting_leaves_the_shared_stage_empty(self): + shared, augmented = build_gpu_stages({"RandomFlipTransformGPU": {"p": 1.0}}, split_shared=False) + self.assertEqual(shared, []) + self.assertEqual([type(t).__name__ for t in augmented], ["RandomFlipTransformGPU"]) + + def test_a_shared_transform_out_of_prefix_position_is_rejected(self): + """Hoisting past a transform that stays behind would silently reorder the pipeline.""" + section = {"ZscoreNormalizationGPU": {"p": 1.0}, "RandomFlipTransformGPU": {"p": 1.0}} + with self.assertRaises(InvalidConfigError) as caught: + build_gpu_stages(section, order_source=OrderSource.CONFIG, split_shared=True) + self.assertIn("RandomFlipTransformGPU", caught.exception.problems[0]) + self.assertIn("prefix", caught.exception.problems[0]) + + def test_the_same_section_still_builds_unsplit(self): + """The prefix rule is the consistency loss's constraint, not the pipeline's.""" + section = {"ZscoreNormalizationGPU": {"p": 1.0}, "RandomFlipTransformGPU": {"p": 1.0}} + built = build_gpu_pipeline(section, order_source=OrderSource.CONFIG) + self.assertEqual([type(t).__name__ for t in built], ["ZscoreNormalizationGPU", "RandomFlipTransformGPU"]) + + def test_the_geometric_group_is_a_prefix_of_pipeline_order(self): + """What makes the split free. `AugType.OCC` exists to keep it true.""" + order = registry.PIPELINE_ORDER[Backend.GPU] + positions = [i for i, name in enumerate(order) if registry.get(name, Backend.GPU).group in SHARED_GROUPS] + self.assertEqual(positions, list(range(len(positions)))) + + +class TestShippedPipelinesAreUnchanged(SmaugLabTestCase): + """A content hash over every shipped GPU pipeline, in every mode. + + `AugType.OCC` moved `RandomCropTransformGPU` out of `AugType.GEO`, and `_buckets` + reads the group to decide what leads the pipeline. Nothing was supposed to move -- + `build.SEQUENTIAL_GROUPS` lists OCC beside GEO for exactly that reason -- and this + is the assertion that says so for all 84 (config, mode) pairs at once. + + Regenerate the hash only when a config or a transform class name is *deliberately* + changed; the failure message prints the new one. + """ + + SIGNATURE = "1a9bfbc31e5e15416c415189dae4aa70c0e9523534a21c1b05c185b37ffcad44" + COUNT = 84 + + def _structures(self) -> dict[str, Any]: + out: dict[str, Any] = {} + for path, config in _buildable_configs(): + for mode in PipelineMode: + out[f"{path.name}:{mode.value}"] = [ + _signature(t) + for t in build_gpu_pipeline( + config.section(Backend.GPU), + mode=mode, + options=config.pipeline_options("random_choose"), + source=config.source, + order_source=config.order_source(), + ) + ] + return out + + def test_the_structure_signature_is_stable(self): + structures = self._structures() + if len(structures) != self.COUNT: + self.skipTest(f"{len(structures)} of {self.COUNT} pipelines built here; the hash is not comparable") + blob = json.dumps(structures, sort_keys=True, separators=(",", ":")) + digest = hashlib.sha256(blob.encode()).hexdigest() + self.assertEqual( + digest, + self.SIGNATURE, + "a shipped GPU pipeline changed shape. If that was intended, set SIGNATURE to " + digest, + ) + + +class TestGpuConfigsAreCoveredBySplitting(SmaugLabTestCase): + def test_every_gpu_config_can_be_split(self): + """No shipped GPU config trips the prefix rule, so none of them blocks the loss.""" + for path in gpu_config_paths(): + if requires_external_asset(path): + continue + config = load_config(str(path)) + with self.subTest(config=path.name): + shared, augmented = build_gpu_stages( + config.section(Backend.GPU), + mode=config.pipeline_mode(), + options=config.pipeline_options("random_choose"), + order_source=config.order_source(), + split_shared=True, + ) + self.assertTrue(shared or augmented) diff --git a/unit_tests/test_cli.py b/unit_tests/test_cli.py index 2501426..f751154 100644 --- a/unit_tests/test_cli.py +++ b/unit_tests/test_cli.py @@ -43,6 +43,12 @@ def test_group_filter(self): self.assertIn("RandomFlipTransformGPU", out) self.assertNotIn("RandomScharrGPU", out) + def test_the_occlusion_group_is_listable(self): + """`occ` is a group like any other, and the crop is the only member.""" + _, out = run("list", "--backend", "gpu", "--group", "occ") + self.assertIn("RandomCropTransformGPU", out) + self.assertNotIn("RandomFlipTransformGPU", out) + def test_json_output_is_machine_readable(self): _, out = run("list", "--backend", "cpu", "--json") payload = json.loads(out) diff --git a/unit_tests/test_config_schema.py b/unit_tests/test_config_schema.py index 72e78c4..3fcf627 100644 --- a/unit_tests/test_config_schema.py +++ b/unit_tests/test_config_schema.py @@ -15,6 +15,7 @@ import tempfile import unittest from pathlib import Path +from typing import ClassVar from smauglab.config import OrderSource, PipelineMode, SmaugConfig, validate_file, validate_section from smauglab.registry import Backend, InvalidConfigError @@ -146,3 +147,100 @@ def test_order_and_mode_are_independent(self): if __name__ == "__main__": unittest.main() + + +class TestLossSection(unittest.TestCase): + """The `loss` section, validated the same way `pipeline` is. + + A typo'd setting has to be an error rather than an ignored key: a config saying + `"detatch_reference": false` would otherwise train with the reference detached and + then be copied into the run folder as that run's only record, claiming the opposite. + """ + + BASE: ClassVar[dict] = {"GPU": {"RandomFlipTransformGPU": {"p": 1.0}}} + + def _config(self, loss) -> SmaugConfig: + return SmaugConfig({**self.BASE, "loss": loss}) + + def test_the_section_is_optional_and_the_term_defaults_to_off(self): + config = SmaugConfig(dict(self.BASE)) + self.assertFalse(config.consistency_enabled()) + self.assertEqual( + config.loss_options("consistency"), + {"enabled": False, "weight": 0.5, "detach_reference": True}, + ) + + def test_defaults_fill_in_around_a_partial_block(self): + config = self._config({"consistency": {"enabled": True, "weight": 0.25}}) + self.assertTrue(config.consistency_enabled()) + self.assertEqual( + config.loss_options("consistency"), + {"enabled": True, "weight": 0.25, "detach_reference": True}, + ) + + def test_options_are_handed_out_as_copies(self): + config = self._config({"consistency": {"enabled": True}}) + config.loss_options("consistency")["weight"] = 99 + self.assertEqual(config.loss_options("consistency")["weight"], 0.5) + + def test_comment_keys_are_accepted(self): + config = self._config({"_why": "ablation", "consistency": {"_note": "x", "enabled": True}}) + self.assertTrue(config.consistency_enabled()) + + def test_an_unknown_term_is_rejected_with_a_suggestion(self): + with self.assertRaises(InvalidConfigError) as caught: + self._config({"consistancy": {"enabled": True}}) + self.assertIn("Did you mean: consistency?", caught.exception.problems[0]) + + def test_an_unknown_setting_is_rejected_with_a_suggestion(self): + with self.assertRaises(InvalidConfigError) as caught: + self._config({"consistency": {"detatch_reference": True}}) + self.assertIn("Did you mean: detach_reference?", caught.exception.problems[0]) + + def test_a_non_boolean_flag_is_rejected(self): + for key in ("enabled", "detach_reference"): + with self.subTest(key=key), self.assertRaises(InvalidConfigError) as caught: + self._config({"consistency": {key: "yes"}}) + self.assertIn(f"loss.consistency.{key}", caught.exception.problems[0]) + + def test_a_negative_weight_is_rejected(self): + with self.assertRaises(InvalidConfigError) as caught: + self._config({"consistency": {"weight": -0.5}}) + self.assertIn("non-negative", caught.exception.problems[0]) + + def test_a_non_numeric_weight_is_rejected(self): + for weight in ("0.5", True, None): + with self.subTest(weight=weight), self.assertRaises(InvalidConfigError): + self._config({"consistency": {"weight": weight}}) + + def test_a_zero_weight_is_accepted(self): + """The "pay the second forward, measure the term, change nothing" ablation.""" + config = self._config({"consistency": {"enabled": True, "weight": 0}}) + self.assertEqual(config.loss_options("consistency")["weight"], 0) + + def test_the_section_must_be_an_object(self): + with self.assertRaises(InvalidConfigError) as caught: + self._config(3) + self.assertEqual(caught.exception.problems, ["loss: expected an object"]) + + def test_a_terms_block_must_be_an_object(self): + with self.assertRaises(InvalidConfigError) as caught: + self._config({"consistency": 3}) + self.assertIn("loss.consistency: expected an object", caught.exception.problems[0]) + + def test_every_problem_in_the_section_is_reported_at_once(self): + with self.assertRaises(InvalidConfigError) as caught: + self._config({"nope": {}, "consistency": {"bad": 1, "weight": -1}}) + self.assertEqual(len(caught.exception.problems), 3) + + def test_a_null_setting_is_rejected_rather_than_meaning_the_default(self): + """`loss_options` merges the block over the defaults, so null would replace one. + + `"weight": null` would reach the trainer as None and raise in `float(...)`; + `"enabled": null` is worse -- `bool(None)` is False, so the term would silently + not run while the recorded config looked like it asked for it. + """ + for key in ("enabled", "weight", "detach_reference"): + with self.subTest(key=key), self.assertRaises(InvalidConfigError) as caught: + self._config({"consistency": {key: None}}) + self.assertIn(f"loss.consistency.{key}", caught.exception.problems[0]) diff --git a/unit_tests/test_consistency_loss.py b/unit_tests/test_consistency_loss.py new file mode 100644 index 0000000..033e0bd --- /dev/null +++ b/unit_tests/test_consistency_loss.py @@ -0,0 +1,376 @@ +"""The augmentation-consistency dice term, and the pipeline split it needs. + +Heavy appearance augmentation teaches the network nothing about the fact that the +augmented and the original patch are the *same anatomy*. This term says so: both go +through the network and their predictions are compared with a soft dice. For that +comparison to mean anything the two images have to be voxel-aligned, which is why the +reference is the output of the geometric stage -- geometry runs once and is shared, so +one segmentation is correct for both branches. + +What is pinned here: the arithmetic of the term, that only the augmented branch gets +gradients, that chaining the two stages reproduces the single container exactly, and +that the guardrails fire on the configs that deserve them. +""" + +from __future__ import annotations + +import importlib.util +import unittest +from typing import ClassVar +from unittest import mock + +import torch + +from smauglab.config import SmaugConfig, load_config +from smauglab.registry import Backend, InvalidConfigError +from smauglab.transforms.build import build_gpu_stages +from unit_tests.helpers import SmaugLabTestCase, configs_dir, seed_everything + +# Raised at import so it works under pytest and `python -m unittest` alike; a +# `pytestmark` would only be understood by one of them. +if importlib.util.find_spec("nnunetv2") is None: + raise unittest.SkipTest("the consistency loss needs the nnunetv2 extra") + +from nnunetv2.training.loss.dice import MemoryEfficientSoftDiceLoss # noqa: E402 + +from smauglab.trainers.nnUNetTrainerDAExt import ( # noqa: E402 + DICE_SMOOTH, + _consistency_problems, + _consistency_warnings, +) +from unit_tests.test_trainers import bare_trainer # noqa: E402 + +SHAPE = (2, 1, 8, 8, 8) + + +def soft_dice(prediction: torch.Tensor, reference: torch.Tensor) -> float: + """Foreground soft dice between two probability maps, written out longhand. + + Deliberately not a call into nnU-Net: this is the independent arithmetic the + trainer's term is checked against. + """ + axes = tuple(range(2, prediction.ndim)) + x, y = prediction[:, 1:], reference[:, 1:] + intersect = (x * y).sum(axes) + dc = (2 * intersect + DICE_SMOOTH) / torch.clip(y.sum(axes) + x.sum(axes) + DICE_SMOOTH, 1e-8) + return float(-dc.mean()) + + +class TwoHeadNet(torch.nn.Module): + """A deep-supervision network: two heads, and a record of each call's grad state.""" + + def __init__(self, *, poison_low_head: bool = False) -> None: + super().__init__() + self.conv = torch.nn.Conv3d(1, 2, 1) + self.poison_low_head = poison_low_head + self.grad_enabled: list[bool] = [] + + def forward(self, x): + self.grad_enabled.append(torch.is_grad_enabled()) + high = self.conv(x) + low = torch.nn.functional.avg_pool3d(high, 2) + if self.poison_low_head: + low = low * float("nan") + return [high, low] + + +def consistency_trainer(*, weight=0.5, detach=True, network=None, batch_dice=False): + """A bare trainer with the term switched on, and identity augmentation stages.""" + trainer = bare_trainer(lambda out, _tgt: (out[0] if isinstance(out, list) else out).mean(), network=network) + trainer.consistency_weight = weight + trainer.consistency_detach = detach + trainer.consistency_nonlin = torch.nn.Softmax(dim=1) + trainer.consistency = MemoryEfficientSoftDiceLoss( + apply_nonlin=trainer.consistency_nonlin, + batch_dice=batch_dice, + do_bg=False, + smooth=DICE_SMOOTH, + ddp=False, + ) + # Identity stages: the split is exercised in TestTwoStagesMatchOneContainer; here + # the point is what train_step does with the reference, not what the pipeline did. + trainer.transforms_shared = lambda data, target: (data, target) + trainer.transforms = lambda data, target: (data, target) + return trainer + + +def batch(key: str = "case_a") -> dict: + torch.manual_seed(7) + return { + "data": torch.rand(SHAPE), + "target": torch.zeros(SHAPE, dtype=torch.long), + "keys": [key, key], + } + + +class TestTheTerm(unittest.TestCase): + def test_it_is_the_soft_dice_of_the_two_predictions(self): + """Checked against the longhand formula, not against another nnU-Net call.""" + network = TwoHeadNet() + trainer = consistency_trainer(network=network) + data = torch.rand(SHAPE) + + with torch.no_grad(): + logits = network(data)[0] + probabilities = torch.softmax(logits, 1) + term = trainer._consistency_term([logits], data) + + # Same input to both branches, so the reference prediction is the same tensor. + self.assertAlmostEqual(float(term), soft_dice(probabilities, probabilities), places=5) + + def test_agreement_does_not_bottom_out_at_minus_one(self): + """A soft target makes this a dice against probabilities, not an agreement score. + + Pinned so nobody later "fixes" a value that is not supposed to be -1: the term's + minimiser over the augmented prediction is the *hard* argmax of the reference, so + it sharpens as well as aligns. Documented on `_consistency_term`. + """ + network = TwoHeadNet() + trainer = consistency_trainer(network=network) + data = torch.rand(SHAPE) + with torch.no_grad(): + logits = network(data)[0] + term = float(trainer._consistency_term([logits], data)) + self.assertGreater(term, -1.0) + self.assertLess(term, 0.0) + + def test_disagreement_scores_worse_than_agreement(self): + network = TwoHeadNet() + trainer = consistency_trainer(network=network) + data = torch.rand(SHAPE) + with torch.no_grad(): + logits = network(data)[0] + + agreeing = float(trainer._consistency_term([logits], data)) + # Flip the two channels: the reference now predicts the opposite class everywhere. + disagreeing = float(trainer._consistency_term([logits.flip(1)], data)) + self.assertGreater(disagreeing, agreeing, "disagreement must cost more than agreement") + + def test_only_the_augmented_branch_gets_gradients_when_detached(self): + network = TwoHeadNet() + trainer = consistency_trainer(network=network, detach=True) + logits = network(torch.rand(SHAPE))[0] + network.grad_enabled.clear() + + trainer._consistency_term([logits], torch.rand(SHAPE)) + self.assertEqual(network.grad_enabled, [False], "the reference forward must run under no_grad") + + def test_the_reference_forward_keeps_the_graph_when_not_detached(self): + network = TwoHeadNet() + trainer = consistency_trainer(network=network, detach=False) + logits = network(torch.rand(SHAPE))[0] + network.grad_enabled.clear() + + trainer._consistency_term([logits], torch.rand(SHAPE)) + self.assertEqual(network.grad_enabled, [True]) + + def test_only_the_full_resolution_head_takes_part(self): + """Head 1 is all NaN; a finite term proves it was never read.""" + network = TwoHeadNet(poison_low_head=True) + trainer = consistency_trainer(network=network) + output = network(torch.rand(SHAPE)) + + term = trainer._consistency_term(output, torch.rand(SHAPE)) + self.assertTrue(torch.isfinite(term), "the deep-supervision heads must not reach the term") + + +class TestTrainStep(unittest.TestCase): + def test_the_term_is_added_with_its_weight(self): + for weight in (0.0, 0.5, 2.0): + with self.subTest(weight=weight): + trainer = consistency_trainer(weight=weight, network=TwoHeadNet()) + trainer.loss = lambda out, _tgt: out[0].mean() * 0.0 + 1.0 + result = trainer.train_step(batch()) + self.assertAlmostEqual( + float(result["loss"]), + 1.0 + weight * result["consistency"], + places=5, + ) + + def test_a_step_reports_the_component(self): + trainer = consistency_trainer(network=TwoHeadNet()) + result = trainer.train_step(batch()) + self.assertIn("consistency", result) + self.assertTrue(torch.isfinite(torch.tensor(result["consistency"]))) + self.assertEqual(trainer._consistency_steps, 1) + + def test_nothing_happens_without_a_shared_stage(self): + """A config with no geometric group leaves the term inert rather than guessing.""" + trainer = consistency_trainer(network=TwoHeadNet()) + trainer.transforms_shared = None + result = trainer.train_step(batch()) + self.assertNotIn("consistency", result) + self.assertEqual(trainer._consistency_steps, 0) + + def test_the_epoch_summary_names_the_contribution(self): + trainer = consistency_trainer(weight=0.5, network=TwoHeadNet()) + for i in range(3): + trainer.train_step(batch(f"case_{i}")) + trainer.logged.clear() + + # Only the subclass's own contribution is under test; the parent hook wants a + # logger and an lr_scheduler this bare instance does not have. + parent = type(trainer).__mro__[1] + with mock.patch.object(parent, "on_train_epoch_end", lambda _self, _outputs: None): + type(trainer).on_train_epoch_end(trainer, []) + + self.assertEqual(len(trainer.logged), 1) + self.assertIn("consistency: soft dice", trainer.logged[0]) + self.assertIn("over 3 batches", trainer.logged[0]) + self.assertIn("weight 0.5", trainer.logged[0]) + + def test_the_epoch_counter_resets(self): + trainer = consistency_trainer(network=TwoHeadNet()) + trainer.train_step(batch()) + parent = type(trainer).__mro__[1] + with mock.patch.object(parent, "on_train_epoch_start", lambda _self: None): + type(trainer).on_train_epoch_start(trainer) + self.assertEqual(trainer._consistency_steps, 0) + self.assertIsNone(trainer._consistency_sum) + + def test_a_non_finite_term_is_blamed_on_the_consistency_loss(self): + """Otherwise the report says verdict=loss, pointing the reader at nnU-Net.""" + trainer = consistency_trainer(network=TwoHeadNet()) + trainer.consistency = lambda _x, _y: torch.tensor(float("nan")) + trainer.train_step(batch()) + + reports = [line for line in trainer.logged if "verdict=" in line] + self.assertEqual(len(reports), 1) + self.assertIn("verdict=consistency", reports[0]) + self.assertIn("consistency=nan", reports[0]) + + +class TestTwoStagesMatchOneContainer(SmaugLabTestCase): + """Chaining the split stages has to reproduce the unsplit container exactly.""" + + CONFIG = "transform_params_gpu.json" + + def _stages(self, *, split: bool): + from smauglab.transforms.gpu.base import AugmentationSequentialCustom + + config = load_config(str(configs_dir() / self.CONFIG)) + shared, augmented = build_gpu_stages( + config.section(Backend.GPU), + mode=config.pipeline_mode(), + options=config.pipeline_options("random_choose"), + order_source=config.order_source(), + split_shared=split, + ) + + def container(*stage): + return AugmentationSequentialCustom( + *stage, + data_keys=["input", "mask"], + same_on_batch=True if config.same_on_batch() else None, + ) + + return (container(*shared) if shared else None), container(*augmented) + + def _inputs(self): + volume = self.tiny_volume() + seg = self.tiny_seg() + return volume, seg + + def test_the_chained_stages_reproduce_the_single_container(self): + volume, seg = self._inputs() + + seed_everything() + _, single = self._stages(split=False) + one_data, one_seg = single(volume.clone(), seg.clone()) + + seed_everything() + geometric, intensity = self._stages(split=True) + self.assertIsNotNone(geometric, "transform_params_gpu.json has an active geometric group") + staged_data, staged_seg = geometric(volume.clone(), seg.clone()) + staged_data, staged_seg = intensity(staged_data, staged_seg) + + self.assertTrue(torch.equal(one_data, staged_data), "splitting the pipeline changed the image") + self.assertTrue(torch.equal(one_seg, staged_seg), "splitting the pipeline changed the mask") + + def test_the_intensity_stage_does_not_write_into_the_reference(self): + """What justifies holding an alias in `train_step` instead of cloning. + + `test_no_inplace_mutation.py` pins the same property, but only for transforms + built with default arguments; this runs the shipped configuration. + """ + volume, seg = self._inputs() + seed_everything() + geometric, intensity = self._stages(split=True) + + reference, reference_seg = geometric(volume, seg) + untouched = reference.clone() + intensity(reference, reference_seg) + + self.assertTrue(torch.equal(reference, untouched), "the intensity stage wrote into the reference") + + def test_the_reference_differs_from_the_augmented_image(self): + volume, seg = self._inputs() + seed_everything() + geometric, intensity = self._stages(split=True) + + reference, reference_seg = geometric(volume, seg) + kept = reference.clone() + augmented, _ = intensity(reference, reference_seg) + + self.assertEqual(kept.shape, augmented.shape) + self.assertFalse(torch.allclose(kept, augmented), "the intensity stage did nothing") + + +class TestGuardrails(unittest.TestCase): + GPU_ONLY: ClassVar[dict] = {"GPU": {"RandomFlipTransformGPU": {"p": 1.0}, "RandomGaussianNoiseGPU": {"p": 1.0}}} + + def test_a_cpu_only_config_is_a_hard_error(self): + """Silently not running would leave the recorded config claiming otherwise.""" + config = load_config(str(configs_dir() / "transform_params.json")) + problems = _consistency_problems(config) + self.assertEqual(len(problems), 1) + self.assertIn("GPU section", problems[0]) + + def test_a_gpu_config_has_no_problems(self): + for name in ("transform_params_gpu.json", "transform_params_hybrid.json"): + with self.subTest(config=name): + self.assertEqual(_consistency_problems(load_config(str(configs_dir() / name))), []) + + def test_an_empty_shared_stage_is_only_a_warning(self): + """transform_params_hybrid.json has flip and affine at p=0; geometry is on the CPU.""" + config = load_config(str(configs_dir() / "transform_params_hybrid.json")) + warnings_out = _consistency_warnings(config) + self.assertTrue(any("shared stage is empty" in w for w in warnings_out)) + + def test_cpu_intensity_augmentations_are_named(self): + config = load_config(str(configs_dir() / "transform_params_hybrid.json")) + warning = next(w for w in _consistency_warnings(config) if "CPU section" in w) + self.assertIn("GaussianNoiseTransform", warning) + # CPU geometry runs before the batch splits, so both branches inherit it and it + # must not be reported. + self.assertNotIn("SpatialTransform", warning) + + def test_the_occlusion_is_reported_with_its_probability(self): + config = load_config(str(configs_dir() / "transform_params_gpu.json")) + warning = next(w for w in _consistency_warnings(config) if "RandomCropTransformGPU" in w) + self.assertIn("p=0.25", warning) + self.assertIn("occluded region", warning) + + def test_an_inactive_occlusion_is_silent(self): + payload = {"GPU": {**self.GPU_ONLY["GPU"], "RandomCropTransformGPU": {"p": 0}}} + warnings_out = _consistency_warnings(SmaugConfig(payload)) + self.assertFalse([w for w in warnings_out if "RandomCropTransformGPU" in w]) + + def test_a_clean_gpu_config_warns_about_nothing(self): + self.assertEqual(_consistency_warnings(SmaugConfig(self.GPU_ONLY)), []) + + +class TestConfigSection(unittest.TestCase): + def test_every_shipped_config_leaves_the_term_off(self): + """Nothing shipped changes behaviour; the term is opt-in.""" + for path in sorted(configs_dir().glob("*.json")): + with self.subTest(config=path.name): + try: + config = load_config(str(path)) + except InvalidConfigError: + continue + self.assertFalse(config.consistency_enabled()) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_registry.py b/unit_tests/test_registry.py index 29d23d1..d68a8eb 100644 --- a/unit_tests/test_registry.py +++ b/unit_tests/test_registry.py @@ -317,6 +317,31 @@ def test_unknown_format_is_rejected(self): registry.render_matrix("yaml") +class TestSharedGeometricGroup(unittest.TestCase): + """What `build_gpu_stages` relies on, asserted against the real registry. + + The consistency loss runs `AugType.GEO` once and shares the result with its + reference branch. That is only sound for transforms that *move* voxels: + `RandomCropTransformGPU` blanks the field of view instead, and run ahead of the + intensity stage its flat fill gets texture painted into it. Hence `AugType.OCC`. + """ + + def test_the_crop_is_not_in_the_geometric_group(self): + self.assertIs(registry.get("RandomCropTransformGPU", Backend.GPU).group, AugType.OCC) + + def test_the_geometric_group_is_exactly_the_voxel_moving_transforms(self): + self.assertEqual( + registry.names(Backend.GPU, group=AugType.GEO), + ["RandomFlipTransformGPU", "RandomAffineGPU"], + ) + + def test_the_geometric_group_leads_the_pipeline(self): + """A prefix, which is what makes the split free of reordering.""" + order = registry.PIPELINE_ORDER[Backend.GPU] + positions = [i for i, name in enumerate(order) if registry.get(name, Backend.GPU).group is AugType.GEO] + self.assertEqual(positions, list(range(len(positions)))) + + class TestModuleHygiene(unittest.TestCase): def test_registry_does_not_import_torch(self): """`smauglab list` must be able to answer without paying for torch.""" diff --git a/unit_tests/test_trainers.py b/unit_tests/test_trainers.py index e1cc8b2..bb93359 100644 --- a/unit_tests/test_trainers.py +++ b/unit_tests/test_trainers.py @@ -37,6 +37,50 @@ raise unittest.SkipTest("the trainer needs the nnunetv2 extra") +def bare_trainer(loss, *, transforms=None, network=None, device="cpu"): + """A trainer instance without plans, dataset_json or an output folder. + + `__init__` wants all three, so `__new__` sidesteps it and only the attributes + `train_step` and the epoch hooks actually read are filled in. Kept at module level + because three test classes need it, and because every attribute the trainer grows + has to be added here or the bare instance raises AttributeError instead of failing + on what is under test. + """ + import torch + + from smauglab.trainers.nnUNetTrainerDAExt import nnUNetTrainerDAExtGPU + + trainer = nnUNetTrainerDAExtGPU.__new__(nnUNetTrainerDAExtGPU) + trainer.device = torch.device(device) # cpu, so autocast is skipped via dummy_context + trainer.transforms = transforms + trainer.transforms_shared = None + trainer.grad_scaler = None + trainer.network = network if network is not None else torch.nn.Conv3d(1, 2, 1) + trainer.optimizer = torch.optim.SGD(trainer.network.parameters(), lr=0.1) + trainer.loss = loss + trainer.current_epoch = 0 + trainer.num_iterations_per_epoch = 250 + trainer._nan_steps_this_epoch = 0 + trainer._nan_steps_total = 0 + # `on_train_epoch_end` reports all three guards; `__init__` sets every counter. + trainer._scale_floor_hits_this_epoch = 0 + trainer._scale_floor_hits_total = 0 + trainer._min_grad_scale = 1.0 + trainer.consistency = None + trainer.consistency_weight = 0.0 + trainer.consistency_detach = True + trainer.consistency_nonlin = torch.nn.functional.softmax + trainer._consistency_sum = None + trainer._consistency_steps = 0 + trainer._consistency_last = None + trainer._get_deep_supervision_scales = lambda: None + # An instance attribute shadowing the bound method: capturing the log this way + # needs no output folder and no open file handle. + trainer.logged = [] + trainer.print_to_log_file = lambda *args, **_kwargs: trainer.logged.append(" ".join(str(a) for a in args)) + return trainer + + def flatten(transform) -> list[str]: """Class names, descending into Compose and unwrapping RandomTransform.""" if isinstance(transform, ComposeTransforms): @@ -265,34 +309,12 @@ class TestNonFiniteLossGuard(unittest.TestCase): Driving this needs a trainer instance, which the rest of this file never builds: `__init__` wants plans, a dataset_json and an output folder to copy the config into. - `__new__` sidesteps all of it, and `train_step` only ever touches seven attributes. + `__new__` sidesteps all of it, and `train_step` only ever touches the handful of + attributes `bare_trainer` sets. """ def _trainer(self, loss, transforms=None): - import torch - - from smauglab.trainers.nnUNetTrainerDAExt import nnUNetTrainerDAExtGPU - - trainer = nnUNetTrainerDAExtGPU.__new__(nnUNetTrainerDAExtGPU) - trainer.device = torch.device("cpu") # so autocast is skipped via dummy_context - trainer.transforms = transforms - trainer.grad_scaler = None - trainer.network = torch.nn.Conv3d(1, 2, 1) - trainer.optimizer = torch.optim.SGD(trainer.network.parameters(), lr=0.1) - trainer.loss = loss - trainer.current_epoch = 0 - trainer.num_iterations_per_epoch = 250 - trainer._nan_steps_this_epoch = 0 - trainer._nan_steps_total = 0 - # `on_train_epoch_end` reports both guards; `__init__` sets all four counters. - trainer._scale_floor_hits_this_epoch = 0 - trainer._scale_floor_hits_total = 0 - trainer._get_deep_supervision_scales = lambda: None - # An instance attribute shadowing the bound method: capturing the log this way - # needs no output folder and no open file handle. - trainer.logged = [] - trainer.print_to_log_file = lambda *args, **_kwargs: trainer.logged.append(" ".join(str(a) for a in args)) - return trainer + return bare_trainer(loss, transforms=transforms) def _batch(self, key="case_a"): import torch @@ -391,19 +413,8 @@ class TestTargetReachesTheDeviceEitherShape(unittest.TestCase): """ def _trainer(self): - import torch - - from smauglab.trainers.nnUNetTrainerDAExt import nnUNetTrainerDAExtGPU - - trainer = nnUNetTrainerDAExtGPU.__new__(nnUNetTrainerDAExtGPU) - trainer.device = torch.device("cpu") - trainer.transforms = None # a CPU-only config builds no GPU pipeline - trainer.grad_scaler = None - trainer.network = torch.nn.Conv3d(1, 2, 1) - trainer.optimizer = torch.optim.SGD(trainer.network.parameters(), lr=0.1) - trainer.current_epoch = 0 - trainer._nan_steps_this_epoch = 0 - trainer._nan_steps_total = 0 + # transforms=None: a CPU-only config builds no GPU pipeline. + trainer = bare_trainer(None) trainer.seen = [] # Record what the loss was handed, which is the thing under test. trainer.loss = lambda out, tgt: (trainer.seen.append(tgt), out.mean())[1]