Repository navigation
Conversation
Integrates the training recipe of the 'mitochondria2' model, which so far only existed outside this repository, as a first-class part of the package. - synapse_net/training/mitochondria.py: get_mitochondria_paths for the recursive data discovery and the train/val split (optionally from a split file), and mitochondria_training, which fixes the hyperparameters of the published model. Registered as the CLI synapse_net.run_mitochondria_training. - supervised_training: expose initial_features, mixed_precision, early_stopping and log_image_interval. The defaults keep the previous behaviour for all other callers. - scripts/cooper/training/train_mitochondria_tomo.py plus the split file of the published run, so it can be reproduced file for file. - test/training/__init__.py, without which unittest discovery skipped all of test/training. Discovery goes from 106 to 121 tests with this commit. The recipe was taken from the init of the deployed checkpoint, not from the training script. Verified by training on the real tomograms and comparing the resulting checkpoint against the deployed one: model kwargs, loss, optimizer, scheduler, mixed_precision, early_stopping, log_image_interval and the loader kwargs all agree, and both models have the same state dict and the same 18 InstanceNorm3d layers. Note that the percentile normalization is part of the model: the same preprocessing has to be passed to segment_mitochondria at inference time. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The previous commit added test/training/__init__.py so that unittest discovery would pick
up that folder. That also enabled the eight tests in test/training/test_training_cli.py,
which had never run in CI before and do not pass there: all three CI jobs failed, while
main is green. Windows failed within seconds, which matches the skip that test/test_cli.py
already carries ("CLI does not work on Windows") for the same kind of subprocess test.
Fixing those tests is not this change's job, so restore test/training to the state it was
in and move the new mitochondria tests to test/, where discovery finds them without
enabling anything else. CI now runs 113 tests, the 106 from main plus these 7.
Also skip the CLI test on Windows, following test/test_cli.py.
The tests in test/training/test_training_cli.py are still invisible to CI. That is
pre-existing and worth a separate look.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Integrates the training recipe of the 'cristae4' model, which so far only existed outside this repository. Cristae training needs machinery that the package did not have, so this adds it rather than only wiring up a recipe. - synapse_net/training/loss.py: MaskedDiceLoss and MaskedDiceLossPerSample, which read a per-voxel loss mask from the second half of the target channels. This is how the voxels of mitochondria without cristae annotations are kept out of the loss. - synapse_net/training/transform.py: the joint transforms that build that mask, the membrane proximity weighting it can carry, and standardize_channel / normalize_channel, which normalize the tomogram channel while leaving the mitochondria state untouched. - synapse_net/training/cristae.py: get_cristae_paths and cristae_training, plus the CLI synapse_net.run_cristae_training. - supervised_training: expose 'norm' (the cristae model has no normalization layers, and until now the model was always built with InstanceNorm) and let the joint 'transform' be overridden, which the cristae mask transform needs. The defaults keep the previous behaviour for all other callers, and the model now records its norm in the checkpoint. - scripts/cooper/training/train_cristae.py plus the split file of the published run. Verified against the originals in the separate synapse repository: the losses (binary and weighted, several shapes), the membrane weighting, both mask transforms including the batched path, and the channel normalizations are all bit-identical. A checkpoint trained with this code agrees with the deployed one on every model kwarg (including norm=None and the two input channels), the optimizer, the scheduler, mixed precision, early stopping, the log interval and the loader kwargs. get_cristae_paths reproduces the published 135 / 14 split file for file, in order. Note that the published checkpoint could not be loaded without the synapse package on the path, because its loss class was pickled into the checkpoint. Checkpoints written by this code no longer depend on anything outside synapse_net. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The previous commit added the eight h5 files that test_cristae_training.py writes to ./tmp_cristae_training, left behind by a test run that did not clean up. setUp creates the datasets in these files in append mode, which fails if they already exist, so the checked-in files broke 18 of the cristae tests on all three CI platforms. Remove them, ignore tmp_*/ so that test output cannot be committed again, and let the new tests clear their tmp folder in setUp, so that an interrupted run does not break the following ones. Windows already failed before that, in the split-file tests. The split entries use '/', and os.path.join(root, "sub/tomo.h5") keeps it, while glob returns the path with '\'. The paths never matched, so besides the failing tests, passing get_cristae_test_paths() as exclude did not keep the test volumes out of a random split on Windows. Join the parts of the entries instead, which gives the same paths on Linux and macOS as before. Do the same in get_mitochondria_paths, and derive the root names with normpath, so that a trailing '\' does not give an empty name. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…nto add-cristae-training Bring in the evaluation scripts for the mitochondria and cristae models (PR #119 and the add-mito-eval-script branch), as they are on these branches. They are adapted to the training recipes of this branch in the following commits. scripts/cooper/run_mitochondria_segmentation.py keeps defaulting to 'mitochondria2'. The mito eval branch predates that model and would otherwise switch the default back to 'mitochondria'. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
If the mitochondria are not passed separately, segment_cristae takes them from the second channel of the input, which then has the dtype of the tomogram. For a float input, e.g. the test volume Otof_AVCN07_455L_KO_M.Stim_B3_2_35933_combined of the cristae split, regionprops raised "Non-integer label_image types are ambiguous" in _erode_instances. Round such a channel to labels first. test_cristae_inference.py runs the post-processing for real, with only the prediction mocked, for float and integer inputs and with erosion. The float tests fail without the fix. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
segment_test_set.py segments the held-out test set of either model: the 11 tomograms in eval_data_h5_s4 for mitochondria, with the percentile normalization of the recipe, and the 'test' entries of the cristae split file. The model is a registry name, a checkpoint folder or a .pt file. '--tile' and '--halo' pin the tiling, which otherwise depends on the GPU memory, and '--save_predictions' also writes pred/foreground and pred/boundary. evaluate_cristae.py is now a port of the evaluation of the published cristae models in the synapse repo (synapse/cristae/evaluate.py and evaluation/cristae/evaluate_cristae_ap.py): voxel-wise dice, precision and recall inside the mitochondria with cristae annotations (state 1), and from the saved predictions the AP and ROC-AUC in the annotated mitochondria and in 8 and 12 nm shells at the membrane and their complements. It replaces the instance F1 and SBD, which are degenerate for the binary cristae segmentation, and the per-volume special case, which the state mask supersedes. evaluate_mitochondria.py defaults to the test set and to the keys segment_test_set.py writes, pairs files by exact name first, and appends the average row instead of overwriting the last dataset. Both scripts print their tables with pandas.to_markdown, so environment.yaml adds tabulate. Checked against the stored outputs of the published runs: on the old segmentations and predictions both scripts reproduce the published scores exactly (mitochondria2 F1 0.8541 / SBD 0.7444, cristae5 dice 0.7581 and AP / band AP), and segment_test_set.py with the published tiling reproduces the stored predictions and cristae segmentations bitwise in the environment they were made with. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…training Both recipes need an 80 GB GPU: the mitochondria recipe trains without mixed precision and needs about 50 GB (it runs out of memory on a 40 GB A100), the cristae recipe needs about 56 GB with mixed precision. '--deterministic' needs more memory still and does not fit on an 80 GB GPU with the default batch sizes, so its help now says that it costs GPU memory as well as throughput. The driver scripts also point to the test-set evaluation, and the mitochondria one to the '--random_split' option it actually has. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
The recipes are now written down once, in synapse_net.training.mitochondria and .cristae, which hard-code the published parameters and keep only the arguments a caller needs to change. The console scripts are removed, together with setup.py's entry points and the argparse duplicates of the recipe, and the driver scripts only hold the data paths, the split files and the run names. Their output folder is now required, so that the defaults cannot overwrite the published checkpoints. supervised_training keeps only what the recipes need: norm, transform, mixed_precision, early_stopping and log_image_interval, and resume, now via torch-em's own load_from_checkpoint. The overwrite guard, which changed the behaviour of every existing caller, and the seeding are removed; '--deterministic' did not fit on an 80 GB GPU with the published batch sizes. loss.py keeps only MaskedDiceLossPerSample, and transform.py merges MitoStateMaskTransform into AugmentedMitoStateMaskTransform and the membrane band helpers into one private function, with unchanged numerics. environment.yaml and setup.py are back to main. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
One script segments the test set of the mitochondria or the cristae model and scores it in memory: instance F1, precision, recall and SBD for mitochondria, and for cristae the voxel dice, precision and recall in the annotated mitochondria plus the AP and ROC-AUC of the prediction there and in 8 and 12 nm membrane bands. It replaces segment_test_set.py, evaluate_mitochondria.py and evaluate_cristae.py, and prints with to_string, so tabulate is not needed. evaluation.py is back to main. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
test_training.py replaces the three training test modules: the data splits (including the Windows paths), the loss numerics, the mask transform, and a short in-process cristae training that is then resumed, instead of two subprocess trainings through the removed CLIs. The fixtures live in a temporary folder, so .gitignore is back to main, and test_cristae_inference.py tests the float fix in one test. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Main moved get_3d_model to training/models.py and get_supervised_loader to training/dataloaders.py. The norm and transform options of this branch now live there, and supervised_training passes norm through _get_model, where it applies to the 3D U-Net without a pretrained backbone. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Main now discovers test/training, so the reason to keep these tests in test/ is gone. Rename the class, which had the same name as the one in test/training/test_training.py, and check that the cristae checkpoint is built without normalization, now that norm goes through _get_model. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The float mitochondria channel in segment_cristae comes from the second input channel, which has the dtype of the tomogram; regionprops only accepts integer labels. evaluate_test_set.py: since #188 the default tiling depends on the model's training patch shape and on the hardware, not only on the GPU memory. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This branch has not been deployed
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
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.
Adds the training recipes of the published
mitochondria2andcristae5models to the package, together with the test-set evaluation, so that both models can be retrained and scored from this repository.Changes
synapse_net.training.mitochondria:get_mitochondria_paths(recursive data discovery, train/val split, optionally from a split file) andmitochondria_trainingwith the hyperparameters ofmitochondria2. The model is trained on percentile-normalized tomograms, so the samepreprocesshas to be passed tosegment_mitochondria.synapse_net.training.cristae:get_cristae_paths,get_cristae_test_pathsandcristae_trainingwith the recipe ofcristae5. The network gets the tomogram and a mitochondria state (0 background, 1 mitochondrion with cristae annotations, 2 without) as two input channels.synapse_net.training.loss.MaskedDiceLossPerSampleandsynapse_net.training.transform.AugmentedMitoStateMaskTransform: keep the mitochondria without cristae annotations out of the loss and weight it towards the membrane of the annotated ones.standardize_channelnormalizes the tomogram channel only.supervised_traininggetsnorm,transform,resume,mixed_precision,early_stoppingandlog_image_interval.normandtransformare passed throughget_3d_model/_get_modelintraining/models.pyandget_supervised_loaderintraining/dataloaders.py. The defaults keep the previous behaviour for all other callers.segment_cristaerounds a float mitochondria channel to labels;regionpropsfailed on it before.scripts/cooper/training: driver scripts and the split files of the published runs, andevaluate_test_set.py, which segments and scores the held-out test set of either model.Verification
mitochondria2andcristae5checkpoints on the model kwargs, loss, optimizer, scheduler, mixed precision, early stopping, log interval and loader kwargs.get_cristae_pathsreproduces the published 135 / 14 split file for file.evaluate_test_set.pyreproduces the published scores (mitochondria2F1 0.8541 / SBD 0.7444,cristae5dice 0.7581 and AP / band AP).test/training/test_mitochondria_cristae_training.py(splits including Windows paths, loss, mask transform, a short cristae training that is resumed and checked fornorm=None) andtest/test_cristae_inference.py.Notes
training.test_training.test_semisupervised_training(DataLoader worker exits), which is unrelated to this PR.🤖 Generated with Claude Code