Skip to content

add training for the mitochondria tomography model - #180

Open
lufre1 wants to merge 1 commit into
mainfrom
add-mitochondria-training
Open

lufre1 wants to merge 1 commit into
mainfrom
add-mitochondria-training

Conversation

@lufre1

@lufre1 lufre1 commented Sep 18, 2026

Copy link
Copy Markdown
Contributor

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.

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant