Hm/fix augmentation correctness - #55
Draft
Hendrik-code wants to merge 2 commits into
Draft
Conversation
… ladders Which augmentations exist, what they accept, and where they sit in the pipeline were encoded four times over: one dispatch ladder in gpu/transforms.py, two in gpu/transforms_list.py, one in cpu/transforms.py, ~900 lines between them. They had already drifted -- the list pipelines passed a `crop=` argument no transform accepts, and ordered SimulateLowRes differently from the sequential one. Each augmentation class now registers itself. The registry (smauglab/registry.py) is the single source of truth for the class a config key maps to, the parameters it accepts (read from the constructor signature, so there is no second schema to drift), its pipeline order, and its GEO/GE/TA group. smauglab/transforms/build.py does the dispatch once for all three pipeline modes. Config handling moves to smauglab/config.py: a sectioned GPU/CPU/pipeline schema, validation that reports every problem in a file at once rather than one per run, and "did you mean" suggestions. A flat, section-less config is now rejected -- the two namespaces overlapped enough that `GaussianBlurTransform` meant different transforms depending on which builder read it. The three nnU-Net trainers collapse into one. Which sections a config populates decides whether augmentation runs on the dataloader, on the batch, or both, so the CPU/GPU/Hybrid split no longer needs a class each. The class keeps the name nnUNetTrainerDAExtGPU: nnU-Net writes it into every checkpoint and resolves the class from it at inference, so renaming it would strand trained models. A `smauglab` CLI answers what exists and whether a config is valid by reading the registry, so it cannot go out of date, and generates the README coverage matrix and the all-augmentations template config that CI checks for staleness. The .gitignore ignores *.json repo-wide to keep per-experiment configs out (see 7ff2088). That silently swallowed three things this change adds and needs tracked -- the generated template, the test fixtures, and the migrator the config error messages point at -- so each gets an explicit un-ignore. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Each of these is silent -- nothing crashes, nothing fails a test, the pipeline just produces something other than what the config asked for. Every fix below has a regression test that fails against the previous implementation. * RandomFlipTransformGPU never read the flip flags its own generator sampled. It recomputed the same `flip_axis`-derived list for every batch element, so it flipped all configured axes, identically, on every call -- three seeded calls gave byte-identical output, and FlipGenerator3D (including its "at least one axis" guarantee) was dead code. * The 1D Gaussian was sampled at `arange(k)` rather than a centred range, putting its peak at index 0. The 3D kernel's maximum sat at corner [0,0,0], so RandomGaussianBlurGPU and RandomUnsharpMaskGPU blurred *and* translated the image about a voxel -- relative to a segmentation mask that is not convolved. * in_seg/out_seg reduced the mask's class axis with `argmax(...) > 0`. For an ordinary single-channel mask argmax over a length-1 axis is always 0, so the result was all-False: in_seg applied the transform nowhere, out_seg applied it everywhere. For a one-hot mask it dropped the first foreground class, since this repository encodes channel c as label c+1. * The single-axis generators drew their "random" axis in make_samplers, which kornia calls once and caches -- the same axis was degraded for a whole training run. CropGenerator3D additionally drew separate axes for the crop and for its position, and neutralised the position to 1.0 (the far edge) using the crop's neutral value instead of 0.5 (centred). * The 2D CPU Scharr x-kernel had [-10, 0, -10] as its middle row, summing to -20 rather than 0. It was not a gradient operator. Predates the registry work. * The elementwise function transforms normalised with a batch-wide min/max, so a volume's augmentation depended on which other volumes shared its batch. * RandomHistogramEqualizationGPU wrote through an `input[:, c]` view, so its non-finite guard `continue`d over values already in the batch. * RandomChooseXTransformsGPU mutated the caller's batch in place, and raised "params must contain 'scale'" for any transform with a kornia parameter generator, because it calls apply_transform directly and skips forward_parameters. RandomLowResTransformGPU also read flags["data_keys"] unguarded, which only the mask path injects. * Blur sigmas and kernel sizes were drawn with Python's `random`, which torch.manual_seed does not reach and which diverges across DDP ranks. They now use smauglab.transforms.rng, built from the unused _shared_rand apparatus that was already sitting in gpu/fromSeg.py. Also: scipy's structuring element rank is taken from the data rather than hardcoded to 3, and `resample_method` is read with .get() so restoring it cannot raise UnboundLocalError. Models trained before this change saw the old behaviour and will not reproduce against it. Configs are unaffected -- no key, parameter or default changed. The README gains a table of what moved and why. `smauglab migrate` was advertised in cli.py, config.py and the README but never existed as a subcommand; those now point at migration/, matching MIGRATE_HINT. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Requires #54
claude: corrected 9 augmentations that apparently did not do what they claimed