feat: augmentation-consistency Dice loss for domain generalization

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:
 

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 <config>`
@@ -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` |

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)
 

smauglab/config.py

@@ -6,7 +6,8 @@
       "_comment": "...",              // any key starting with '_' is ignored
       "GPU":      { "<ClassName>": { <constructor kwargs> }, ... },
       "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 = "<config>
 
 
 #: 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:

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):