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) <[email protected]>
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) <[email protected]>
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) <[email protected]>
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) <[email protected]>
…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) <[email protected]>
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) <[email protected]>
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) <[email protected]>
…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) <[email protected]>
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) <[email protected]>
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) <[email protected]>
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) <[email protected]>
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.
No description provided.