From 8a429bec14e143f25051230d965a8c7bf98620a4 Mon Sep 17 00:00:00 2001 From: iback Date: Mon, 12 Oct 2026 13:27:02 +0000 Subject: [PATCH] feat: compare the augmented and unaugmented prediction with a soft Dice Heavy appearance augmentation never tells the network that the augmented and the original patch are the same anatomy and should get the same segmentation. This adds an optional second loss term that does: both go through the network and their predictions are compared with a soft Dice. The reference is the output of the geometric stage, not the dataloader's patch. The geometric transforms are the only ones that move voxels, so running them once and sharing the result is what makes the two images voxel-aligned and lets one segmentation be correct for both branches. `build_gpu_stages` does the split; `build_gpu_pipeline` is now a thin wrapper over it, so the unsplit path cannot diverge, and `_split_shared` refuses to hoist anything that is not already a prefix of the pipeline rather than silently reordering it. `RandomCropTransformGPU` moves from `AugType.GEO` to a new `AugType.OCC`. It occludes the field of view instead of moving voxels, and `PIPELINE_ORDER` puts it near the end for a reason: in front of the intensity stage, the flat slab it leaves behind gets texture painted into it, which is the failure its own `pad_value` docstring describes. `build.SEQUENTIAL_GROUPS` lists OCC beside GEO so nothing about the built pipelines changes -- `test_builder.TestShippedPipelinesAreUnchanged` hashes all 84 (config, mode) pairs to prove it -- while `build.SHARED_GROUPS` leaves it out of the shared stage. Configured by a new top-level `loss` section, off by default, validated like `pipeline` is. A null setting is rejected rather than read as "use the default": `loss_options` merges the block over `LOSS_DEFAULTS`, so `"weight": null` would reach the trainer as None and `"enabled": null` would quietly turn the term off while the recorded config looked like it asked for it. The term reuses `MemoryEfficientSoftDiceLoss`, not `SoftDiceLoss`: on nnunetv2 2.6 `get_tp_fp_fn_tn` builds `y_onehot` as a bool tensor and computes `~y_onehot`, which is a TypeError on the soft target this needs. Only the full-resolution head takes part, and the epoch summary prints the component on its own line because `train_loss` now carries both terms. Known interaction, warned about at startup and documented in the README: the crop is not applied to the reference, so in the batches where it fires the term asks the network to agree with the reference inside the occluded region, where the supervised target says background. Nothing masks it; `p: 0` avoids it. Verified: full suite in two environments with no new failures, and an end-to-end run of the real trainer on Dataset400_nose_nosemain confirming the total is the supervised loss plus weight times the term, and that `weight: 0` measures it while changing nothing. Co-Authored-By: Claude Opus 5 --- CONTRIBUTING.md | 1 + README.md | 57 +++- smauglab/cli.py | 2 +- smauglab/config.py | 111 ++++++- smauglab/registry.py | 16 +- smauglab/trainers/nnUNetTrainerDAExt.py | 232 +++++++++++++-- smauglab/transforms/build.py | 89 +++++- smauglab/transforms/gpu/spatial.py | 7 +- unit_tests/test_builder.py | 183 +++++++++++- unit_tests/test_cli.py | 6 + unit_tests/test_config_schema.py | 98 ++++++ unit_tests/test_consistency_loss.py | 376 ++++++++++++++++++++++++ unit_tests/test_registry.py | 25 ++ unit_tests/test_trainers.py | 80 ++--- 14 files changed, 1208 insertions(+), 75 deletions(-) create mode 100644 unit_tests/test_consistency_loss.py 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 cb2b381..18698ee 100644 --- a/README.md +++ b/README.md @@ -168,6 +168,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 ` @@ -215,7 +270,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 3126390..fbd331d 100644 --- a/smauglab/trainers/nnUNetTrainerDAExt.py +++ b/smauglab/trainers/nnUNetTrainerDAExt.py @@ -31,15 +31,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. @@ -57,6 +58,11 @@ #: warning worth reading becomes noise nobody reads. NAN_REPORTS_PER_EPOCH = 3 +#: 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: """Locate the config: new env var, then the deprecated one, then the default.""" @@ -110,6 +116,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. @@ -123,28 +186,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")) @@ -154,6 +272,12 @@ def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dic self._nan_steps_this_epoch = 0 self._nan_steps_total = 0 + # 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, ...]], @@ -266,6 +390,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) @@ -277,6 +411,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() @@ -302,8 +438,51 @@ 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 _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. @@ -336,17 +515,25 @@ 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._consistency_sum = None + self._consistency_steps = 0 def on_train_epoch_end(self, train_outputs: list[dict]): super().on_train_epoch_end(train_outputs) @@ -359,3 +546,14 @@ def on_train_epoch_end(self, train_outputs: list[dict]): f"this epoch ({self._nan_steps_total} since the run started); " f"those steps did not update the weights" ) + 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 c3e44d4..33e787f 100644 --- a/unit_tests/test_trainers.py +++ b/unit_tests/test_trainers.py @@ -37,6 +37,46 @@ 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 + 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,31 +305,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 - 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 @@ -388,19 +409,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]