From b2a822a59610a10fcd7a6a1de78bbd5cf0bdad49 Mon Sep 17 00:00:00 2001 From: Suraj Mishra Date: Tue, 10 Nov 2026 15:44:59 +0000 Subject: [PATCH] Make the paper's inference and training settings the defaults predict: 30% window overlap and constant (uniform) blending instead of 50% and Gaussian; a voxel is foreground at probability >= threshold so the mask equals the saved probability map thresholded at 128. train/finetune: batch size 32, up to 1000 epochs, 20% validation hold-out; fine-tuning uses lr 5e-5 and patience 50, from scratch keeps 1e-4 and 100. All values are from the paper's Methods and the cyto-dl configs used for its experiments. Docs, changelog and README benchmark numbers updated; a test pins the defaults. Assisted-by: claude-code (claude-fable-5-1) Co-Authored-By: Claude Fable 5.1 --- CHANGELOG.md | 11 +++++++++++ README.md | 17 ++++++++++------- docs/SEGMENT.md | 18 ++++++++++++------ docs/TRAINING.md | 32 ++++++++++++++++++-------------- src/subcellseg/cli.py | 32 ++++++++++++++++++++------------ src/subcellseg/data.py | 2 +- src/subcellseg/inference.py | 18 ++++++++++++------ src/subcellseg/training.py | 24 ++++++++++++++++++------ tests/test_subcellseg.py | 22 +++++++++++++++++++++- 9 files changed, 123 insertions(+), 53 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index db62a84..e744cdb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,17 @@ All notable changes to SubCellSeg are listed here. The format follows ## Unreleased +### Changed +- `predict` defaults now match the paper: 30% window overlap and uniform + (`constant`) blending, instead of 50% and Gaussian. Pass + `--overlap 0.5 --blend gaussian` for the previous behaviour. A voxel is + foreground when its probability is at least the threshold (was: above), so + `_seg.tif` is exactly `_prob.tif >= 128`. +- `train` and `finetune` defaults now match the paper: batch size 32 (was 8), + at most 1000 epochs (was 700), 20% of the images held out for validation + (was 10%). Fine-tuning uses a learning rate of 5e-5 and stops after 50 + epochs without improvement; training from scratch keeps 1e-4 and 100. + ### Added - Benchmarks section in the README and `tools/benchmark/` scripts. - Contributing guide, code of conduct, security policy, citation file, diff --git a/README.md b/README.md index 6536f30..42302ef 100644 --- a/README.md +++ b/README.md @@ -144,18 +144,21 @@ from `https://animatedcell-test-data.s3.us-west-2.amazonaws.com/`. Wall time is the whole command: starting Python and PyTorch, loading the model, reading the image over the network (1 to 5 s here) and writing the outputs. -The sliding-window inference itself takes about 1.3 s for image 40 on the A100 -(63 windows at roughly 48 windows/s). The default `--batch-size 4` and -`--patch-size 32 128 128` were used. +The sliding-window inference itself took about 1.3 s for image 40 on the A100. +Timings were measured with the defaults of the time, 50% window overlap and +Gaussian blending (252 windows for image 40, in batches of 4). The current +defaults are the paper's settings, 30% overlap and constant blending, which +run 140 windows for the same image and are a little faster. The CPU and GPU masks for image 40 agree on 99.997 % of voxels. Against the -classic `SEG_STRUCT` segmentation stored in the same files: +classic `SEG_STRUCT` segmentation stored in the same files, with the current +defaults (measured 2026-11-10): | Image | Dice | IoU | Precision | Recall | |---|---|---|---|---| -| 40 | 0.533 | 0.363 | 0.732 | 0.419 | -| 6919 | 0.877 | 0.781 | 0.902 | 0.853 | -| 11370 | 0.846 | 0.733 | 0.883 | 0.812 | +| 40 | 0.532 | 0.362 | 0.733 | 0.417 | +| 6919 | 0.874 | 0.777 | 0.904 | 0.847 | +| 11370 | 0.845 | 0.732 | 0.888 | 0.806 | To reproduce, run `tools/benchmark/bench_predict.sh cuda results 40 6919 11370` (it times `subcellseg predict`, records GPU memory, and runs diff --git a/docs/SEGMENT.md b/docs/SEGMENT.md index 4a3c7ea..cebaec0 100644 --- a/docs/SEGMENT.md +++ b/docs/SEGMENT.md @@ -115,12 +115,14 @@ Run `subcellseg predict --help` to see them. The ones you are most likely to use | `--input`, `-i` | Image, folder of images, or web address | required | | `--output`, `-o` | Folder for the results | `subcellseg_output` | | `--channel` | Channel with the structure (number from 0, or name) | `0` | -| `--threshold` | Score above which a voxel counts as structure | `0.5` | +| `--threshold` | Score at or above which a voxel counts as structure | `0.5` | | `--save-prob` | Also save the score image | off | | `--checkpoint`, `-c` | Use your own model | pretrained model | | `--overwrite` | Redo images that are already done | off | | `--device` | `auto`, `cpu` or `cuda` | `auto` | | `--batch-size` | Lower it if the GPU runs out of memory | `4` | +| `--overlap` | How much neighbouring windows overlap | `0.3` | +| `--blend` | How overlapping windows are combined: `constant` or `gaussian` | `constant` | ## Problems @@ -145,12 +147,16 @@ images. ## How this relates to the paper -`subcellseg predict` blends overlapping windows with 50% overlap and Gaussian -weighting, which gives smoother results at window borders. The benchmarks in -the paper used 30% overlap and uniform weighting. To reproduce those settings: +The defaults of `subcellseg predict` are the settings behind the numbers in +the paper: windows of 32 x 128 x 128 voxels, 30% overlap between neighbouring +windows, a uniform (`constant`) average where windows overlap, and a threshold +of 0.5 on the probability map (128 in the saved `_prob.tif`). + +Gaussian blending gives smoother transitions at window borders. It is not what +the paper used, and the masks differ slightly: ``` -subcellseg predict --input my_images --output results --overlap 0.3 --blend constant +subcellseg predict --input my_images --output results --overlap 0.5 --blend gaussian ``` - + diff --git a/docs/TRAINING.md b/docs/TRAINING.md index a0ae8ba..03125b1 100644 --- a/docs/TRAINING.md +++ b/docs/TRAINING.md @@ -103,7 +103,7 @@ SubCellSeg sets a few of your examples aside for checking its progress (the "validation" images) and learns from the rest. You will see lines like: ``` -Training on 11 image(s), validating on 1. +Training on 10 image(s), validating on 2. Starting from weights in C:\Users\you\.cache\subcellseg\subcellseg_v1.pt Training on the GPU. Epoch 12: 100%|##########| 2/2 [00:01<00:00, 1.70it/s, val/loss=0.351, val/dice=0.693, train/loss=0.471] @@ -121,7 +121,7 @@ The numbers jump around from one epoch to the next, especially with few images. Look at the trend, not at single values. Training stops by itself when the validation images have stopped improving -for 100 epochs, or after 700 epochs. You can also stop it at any time with +for 50 epochs (100 when training from scratch), or after 1000 epochs. You can also stop it at any time with `Ctrl+C`; the best model so far is already saved. When it finishes it tells you where the model is: @@ -178,9 +178,10 @@ Run `subcellseg finetune --help` to see all of them. | `--train-csv` | The list from Step 3 | required | | `--output`, `-o` | Folder for the new model | `subcellseg_run` | | `--channel` | Channel with the structure, if images have several (number from 0, or name) | `0` | -| `--epochs` | Stop after this many epochs at the latest | `700` | -| `--patience` | Stop when there was no improvement for this many epochs | `100` | -| `--batch-size` | Lower it if the GPU runs out of memory | `8` | +| `--epochs` | Stop after this many epochs at the latest | `1000` | +| `--patience` | Stop when there was no improvement for this many epochs | `50` (fine-tune), `100` (from scratch) | +| `--batch-size` | Lower it if the GPU runs out of memory | `32` | +| `--lr` | Learning rate | `5e-5` (fine-tune), `1e-4` (from scratch) | | `--val-csv` | Your own list of validation images, instead of an automatic choice | automatic | | `--cache-dir` | Folder for prepared images, for speed | off | | `--resume` | Continue an interrupted run from `last.ckpt` | off | @@ -191,18 +192,20 @@ Run `subcellseg finetune --help` to see all of them. |---|---| | 2 | about 1.5 GB | | 4 | about 3 GB | -| 8 (default) | about 5 GB | +| 8 | about 5 GB | | 16 | about 10 GB | +| 32 (default) | about 20 GB | -A larger batch size makes each step more reliable but is not required for -good results. The models in the paper were trained with a batch size of 32. +The default is the batch size the models in the paper were trained with. A +smaller batch size also gives good results and is the first thing to lower +when the GPU runs out of memory. ## Problems **Training is extremely slow on Windows.** The batch size is probably too big for your graphics card. Windows does not stop with an error when GPU memory runs out; it quietly switches to much slower memory instead. Stop with -`Ctrl+C` and run again with a smaller `--batch-size`, for example 4 or 2. +`Ctrl+C` and run again with a smaller `--batch-size`, for example 8 or 4. **"CUDA out of memory".** Use a smaller `--batch-size`. @@ -243,12 +246,13 @@ subcellseg train --train-csv my_data/train.csv --output my_model - Every epoch, one random 3D patch of 32 x 128 x 128 voxels (Z, Y, X) is cut from each training image, with mild random brightness and contrast changes. - The network is a 3D residual U-Net (MONAI `DynUNet`), trained with the - Generalized Dice + Focal loss, the Adam optimiser (learning rate 0.0001, - weight decay 0.0001) and a learning rate that decays by 0.5% per epoch. + Generalized Dice + Focal loss, the Adam optimiser (learning rate 0.00005 + when fine-tuning, 0.0001 from scratch; weight decay 0.0001) and a learning + rate that decays by 0.5% per epoch. - Validation uses four fixed patches per validation image, the same ones every epoch, so that the numbers are comparable between epochs. -These are the settings used for the models in the paper, except for the fixed -validation patches and the default batch size. +These are the settings used for the models in the paper, with one difference: +the paper's runs cut a fresh random validation patch every epoch. - + diff --git a/src/subcellseg/cli.py b/src/subcellseg/cli.py index 3e74b72..08f96a0 100644 --- a/src/subcellseg/cli.py +++ b/src/subcellseg/cli.py @@ -38,7 +38,7 @@ def _add_predict(sub) -> None: "--threshold", type=float, default=0.5, - help="Probability above which a voxel is foreground (default: 0.5).", + help="Probability at or above which a voxel is foreground (default: 0.5).", ) p.add_argument( "--save-prob", @@ -47,13 +47,17 @@ def _add_predict(sub) -> None: ) p.add_argument("--device", default="auto", help="auto, cpu, cuda, cuda:1, ... (default: auto).") p.add_argument( - "--overlap", type=float, default=0.5, help="Sliding-window overlap (default: 0.5)." + "--overlap", + type=float, + default=0.3, + help="Fraction by which neighbouring windows overlap (default: 0.3, as in the paper).", ) p.add_argument( "--blend", choices=["gaussian", "constant"], - default="gaussian", - help="How overlapping windows are combined (default: gaussian).", + default="constant", + help="How overlapping windows are combined: constant averages them uniformly " + "(default, as in the paper), gaussian favours the centre of each window.", ) p.add_argument( "--patch-size", @@ -122,21 +126,21 @@ def _add_train(sub, name: str, finetune: bool) -> None: p.add_argument( "--val-fraction", type=float, - default=0.1, - help="Fraction held out for validation without --val-csv (default: 0.1).", + default=0.2, + help="Fraction held out for validation without --val-csv (default: 0.2).", ) - p.add_argument("--epochs", type=int, default=700, help="Maximum epochs (default: 700).") + p.add_argument("--epochs", type=int, default=1000, help="Maximum epochs (default: 1000).") p.add_argument( "--patience", type=int, - default=100, - help="Stop after this many epochs without improvement (default: 100).", + help="Stop after this many epochs without improvement " + "(default: 50 when fine-tuning, 100 when training from scratch).", ) p.add_argument( "--batch-size", type=int, - default=8, - help="Patches per step; lower it if the GPU runs out of memory (default: 8).", + default=32, + help="Patches per step; lower it if the GPU runs out of memory (default: 32).", ) p.add_argument( "--patches-per-image", @@ -152,7 +156,11 @@ def _add_train(sub, name: str, finetune: bool) -> None: metavar=("Z", "Y", "X"), help="Training patch size (default: 32 128 128).", ) - p.add_argument("--lr", type=float, default=1e-4, help="Learning rate (default: 1e-4).") + p.add_argument( + "--lr", + type=float, + help="Learning rate (default: 5e-5 when fine-tuning, 1e-4 when training from scratch).", + ) p.add_argument( "--num-workers", type=int, default=4, help="Data loading processes (default: 4)." ) diff --git a/src/subcellseg/data.py b/src/subcellseg/data.py index b4880e8..6c2f0f0 100644 --- a/src/subcellseg/data.py +++ b/src/subcellseg/data.py @@ -56,7 +56,7 @@ def read_manifest(csv_path: PathLike, raw_col: str = "raw", seg_col: str = "seg" def split_items( - items: list[Item], val_fraction: float = 0.1, seed: int = 12345 + items: list[Item], val_fraction: float = 0.2, seed: int = 12345 ) -> tuple[list[Item], list[Item]]: """Randomly hold out a validation set. At least one image goes to each side.""" if len(items) < 2: diff --git a/src/subcellseg/inference.py b/src/subcellseg/inference.py index 32a1372..0936713 100644 --- a/src/subcellseg/inference.py +++ b/src/subcellseg/inference.py @@ -71,8 +71,8 @@ def predict_proba( volume: np.ndarray, *, patch_size: Sequence[int] = PATCH_SIZE, - overlap: float = 0.5, - blend: str = "gaussian", + overlap: float = 0.3, + blend: str = "constant", batch_size: int = 4, progress: bool = False, ) -> np.ndarray: @@ -88,7 +88,11 @@ def predict_proba( overlap Fraction by which neighbouring windows overlap. blend - How overlapping windows are combined: "gaussian" or "constant". + How overlapping windows are combined: "constant" averages them + uniformly, "gaussian" weights the centre of each window more. + + The defaults (30% overlap, constant blending) are the settings behind + the results in the paper. batch_size Number of windows run through the network at once. """ @@ -123,13 +127,14 @@ def predict_proba( def segment(self, volume: np.ndarray, threshold: float = 0.5, **kwargs) -> np.ndarray: """Binary mask of a (Z, Y, X) volume: uint8, 255 inside the structure. - Extra keyword arguments are passed to :meth:`predict_proba`. + A voxel is inside the structure when its probability is at least + ``threshold``. Extra keyword arguments are passed to :meth:`predict_proba`. """ return to_mask(self.predict_proba(volume, **kwargs), threshold) def to_mask(probability: np.ndarray, threshold: float = 0.5) -> np.ndarray: - return (probability > threshold).astype(np.uint8) * 255 + return (probability >= threshold).astype(np.uint8) * 255 def predict( @@ -162,7 +167,8 @@ def predict( channel Channel to segment, as an index or a channel name. threshold - Probability above which a voxel is foreground. + Probability at or above which a voxel is foreground. With the default + of 0.5 the mask equals the saved probability map thresholded at 128. save_probability Also write the probability map as ``name_prob.tif`` (uint8, 0-255). device diff --git a/src/subcellseg/training.py b/src/subcellseg/training.py index e92615f..12d914a 100644 --- a/src/subcellseg/training.py +++ b/src/subcellseg/training.py @@ -91,13 +91,13 @@ def train( channel: int | str = 0, raw_col: str = "raw", seg_col: str = "seg", - val_fraction: float = 0.1, + val_fraction: float = 0.2, patch_size: Sequence[int] = PATCH_SIZE, - batch_size: int = 8, + batch_size: int = 32, patches_per_image: int = 1, - max_epochs: int = 700, - patience: int = 100, - lr: float = 1e-4, + max_epochs: int = 1000, + patience: int | None = None, + lr: float | None = None, weight_decay: float = 1e-4, lr_gamma: float = 0.995, num_workers: int = 4, @@ -133,7 +133,13 @@ def train( Random patches drawn from each image per epoch. max_epochs, patience Train for at most ``max_epochs``; stop early once the validation loss - has not improved for ``patience`` epochs. + has not improved for ``patience`` epochs. ``patience`` defaults to 50 + when fine-tuning and 100 when training from scratch. + lr + Learning rate. Defaults to 5e-5 when fine-tuning and 1e-4 when + training from scratch. + + The defaults are the settings used for the models in the paper. cache_dir Directory for an on-disk cache of preprocessed images. Speeds up training considerably when images are large. @@ -154,6 +160,11 @@ def train( _quiet_lightning() pl.seed_everything(seed, workers=True) output_dir = Path(output_dir) + finetuning = init is not None + if lr is None: + lr = 5e-5 if finetuning else 1e-4 + if patience is None: + patience = 50 if finetuning else 100 train_items = read_manifest(train_csv, raw_col, seg_col) if val_csv is not None: @@ -186,6 +197,7 @@ def train( print(f"Starting from weights in {source}") elif resume is None: print("Training from scratch.") + print(f"Learning rate {lr:g}; stopping after {patience} epochs without improvement.") use_gpu = device in ("gpu", "cuda") or (device == "auto" and torch.cuda.is_available()) if use_gpu: diff --git a/tests/test_subcellseg.py b/tests/test_subcellseg.py index e3bd4ea..c40c595 100644 --- a/tests/test_subcellseg.py +++ b/tests/test_subcellseg.py @@ -71,12 +71,32 @@ def test_predict_proba(model): mask = model.segment(volume, patch_size=SMALL_PATCH) assert mask.dtype == np.uint8 - assert np.array_equal(mask > 0, prob > 0.5) + assert np.array_equal(mask > 0, prob >= 0.5) with pytest.raises(ValueError): model.predict_proba(volume[0]) +def test_defaults_match_paper(): + """The package defaults are the inference and training settings of the paper.""" + import inspect + + from subcellseg.training import train + + infer = inspect.signature(SubCellSeg.predict_proba).parameters + assert infer["patch_size"].default == (32, 128, 128) + assert infer["overlap"].default == 0.3 + assert infer["blend"].default == "constant" + assert inspect.signature(SubCellSeg.segment).parameters["threshold"].default == 0.5 + + fit = inspect.signature(train).parameters + assert fit["batch_size"].default == 32 + assert fit["max_epochs"].default == 1000 + assert fit["val_fraction"].default == 0.2 + # Learning rate and patience depend on whether we fine-tune; resolved inside train(). + assert fit["lr"].default is None and fit["patience"].default is None + + def test_predict_ignores_intensity_scale(model): volume = _volume().astype(np.float32) a = model.predict_proba(volume, patch_size=SMALL_PATCH)