[stmartineau99]

Initial PR

Goal

• Integrating Marei's work from micromatch into synapse-net
• add support for new models (e.g., cryosiam)

Reorganization

• training/dataloaders.py - moved all stuff related to dataloaders here
• training/models.py - moved all stuff related to models here, mostly copied from micromatch/util.py, also the preexisting get_2d_model and get_3d_model

Changes
semisupervised_learning.py

• backbone and model_type - currently used to select UNETR backbone x model type (needs rework when other backbones are added)
• added teacher warmup with teacher_warmup_iterations
• semisuperivsed_training is now just a wrapper around supervised_training, which performs the warmup, and mean_teacher_adaption

domain_adaptation.py

• when backbone is provided, get_raw_transform selects the appropriate raw_transform for that backbone
• get_raw_transform returns clip_max which is used to ensure the augmentations stay in the correct range
• ensured that UNETR model is loaded properly from the warmup checkpoint using torch_em.util.load_model; previously loading from the checkpoint without passing model failed due to UNETR Class not setting init_kwargs (though I fixed this yesterday in torch-em with Anwai)

Deprecated

• preexisting semisupervised_training renamed to semisupervised_training_v0 - I think nobody used this so it can be removed
• weak_augmentations not used anymore since with switched to the new mean teacher setup - but I kept it around in dataloaders.py just in case

Next Steps

• extend training/models.py with cryo-specific models from the torch-em "CryoZoo"

Response to Comments

• removed support for micro_sam.v2 for now

[constantinpape]

Thanks for this PR! The integration of the micromatch backbones is a good step. Please address all inline comments. These are some bugs (two of them I could reproduce), redundant code paths, and smaller issues.

General issues without an inline anchor

• Please bump the torch-em requirement in setup.py and environment.yaml to >=0.10.8 (see the comment on separate_backward).
• flake8 --max-line-length 120 reports about 25 issues in the changed files: missing blank lines (E302), trailing whitespace, no newline at the end of models.py, and lines that are too long. Please fix them.
• No test covers the backbone path or the new semisupervised_training. A small test with a tiny 2D patch and check=True (no download) would catch problems such as the raw_key issue.

Code that should go into torch-em
Most of models.py is general ViT-UNETR code and not specific to synapses. torch-em already contains the related code: preprocess_vit_inputs, the backbone handling in UNETR, and pooch-based checkpoint download in the new cryozoo.py. I suggest that we move these parts upstream:

Part	Reason
get_unetr_model, _get_vit_type	A small factory for UNETR2D/UNETR3D. micromatch and synapse-net can then share it, and we fix the embed_dim bug in one location.
_get_dinov2_checkpoint, _get_dinov3_checkpoint, _get_microsam2_checkpoint, _get_checkpoint	Weight registries and download belong next to the model code (same pattern as CryoZoo.get_checkpoint). The microSAM2 entry probably belongs in micro_sam.
get_raw_transform(backbone) / clip_max	preprocess_vit_inputs already defines the expected input range for each backbone. A get_vit_raw_transform(backbone) next to it keeps the two consistent.
_init_microsam_decoder	General decoder initialization from a pretrained checkpoint. micro_sam's get_unetr has similar logic.
Weak augmentation parameters	MeanTeacherAugmenters should accept an augmentation dict, so that we do not need to change the torch-em module global.
get_unsupervised_loader (optional)	A default_raw_loader with sample masks in torch-em could replace the manual ConcatDataset.

The U-Net factories, get_supervised_loader (boundary and ignore-label setup) and the warmup + mean-teacher logic in semisupervised_training should stay in synapse-net.

[constantinpape]

Above is Claude's review, but I checked all it's findings and they make sense and should be addressed. I would first refactor the general-purpose things to torch-em and create a PR for that there.

[stmartineau99]

Bug: decoder_path is only set for sam and for sam2 with hvit_t_em_organelles. For dinov2, dinov3 and the other sam2 models, return_decoder_path=True raises UnboundLocalError. Thus get_unetr_model(..., init_decoder=True) crashes for a 2D dino model, and the decoder_path is None warning in _init_microsam_decoder cannot occur. Please set decoder_path = None for these backbones.

Done in b7b2bd2.

[constantinpape]

Thanks for the updates! Most of the review points are fixed now. Please see the new inline comment about aug_dict.

Tests
Please add these tests in test/training (tiny synthetic data, few iterations, the same setup as test_training_cli.py):

• semisupervised_training with the U-Net: check that it creates the warmup checkpoint and the final checkpoint, and that a second run reuses the warmup checkpoint. Use a raw_key other than "raw".
• get_raw_transform: the output range for "sam" ([0, 255]) and "dinov2"/"dinov3" ([0, 1]), (None, None) for None, and a ValueError for an unknown backbone.
• One smoke test for the backbone path: supervised_training with a backbone for a few iterations on a 2D patch. If the weight download is too large for CI, skip this test there.

[constantinpape]

The aug_dict argument of MeanTeacherAugmenters is not in torch-em yet (main / 0.10.8 has only ndim, teacher, student, clip_max). Until it is released, this call raises a TypeError for every call of mean_teacher_adaptation. This includes the default synapse_net.run_domain_adaptation CLI and semisupervised_training.

Please make this change part of the new torch-em PR. When that version is released, update the torch-em pin in setup.py and environment.yaml to it.

[stmartineau99]

I meant to leave this PR open until after the torch_em one is done, I will make a PR after the retreat.

[constantinpape]

In addition to the review above: I want to wait with merging this till the PR that moves the shared functionality to torch_em is there, which should also include the required changes for aug_dict (see comment above).

[stmartineau99]

Thanks for this PR! The integration of the micromatch backbones is a good step. Please address all inline comments. These are some bugs (two of them I could reproduce), redundant code paths, and smaller issues.

General issues without an inline anchor

• Please bump the torch-em requirement in setup.py and environment.yaml to >=0.10.8 (see the comment on separate_backward).
• flake8 --max-line-length 120 reports about 25 issues in the changed files: missing blank lines (E302), trailing whitespace, no newline at the end of models.py, and lines that are too long. Please fix them.
• No test covers the backbone path or the new semisupervised_training. A small test with a tiny 2D patch and check=True (no download) would catch problems such as the raw_key issue.

Code that should go into torch-em Most of models.py is general ViT-UNETR code and not specific to synapses. torch-em already contains the related code: preprocess_vit_inputs, the backbone handling in UNETR, and pooch-based checkpoint download in the new cryozoo.py. I suggest that we move these parts upstream:

Part Reason
get_unetr_model, _get_vit_type A small factory for UNETR2D/UNETR3D. micromatch and synapse-net can then share it, and we fix the embed_dim bug in one location.
_get_dinov2_checkpoint, _get_dinov3_checkpoint, _get_microsam2_checkpoint, _get_checkpoint Weight registries and download belong next to the model code (same pattern as CryoZoo.get_checkpoint). The microSAM2 entry probably belongs in micro_sam.
get_raw_transform(backbone) / clip_max preprocess_vit_inputs already defines the expected input range for each backbone. A get_vit_raw_transform(backbone) next to it keeps the two consistent.
_init_microsam_decoder General decoder initialization from a pretrained checkpoint. micro_sam's get_unetr has similar logic.
Weak augmentation parameters MeanTeacherAugmenters should accept an augmentation dict, so that we do not need to change the torch-em module global.
get_unsupervised_loader (optional) A default_raw_loader with sample masks in torch-em could replace the manual ConcatDataset.
The U-Net factories, get_supervised_loader (boundary and ignore-label setup) and the warmup + mean-teacher logic in semisupervised_training should stay in synapse-net.

I am working on the new torch-em PR:

I have a comment about the generic decoder initialization (init_decoder), I'd like to add a new method to UNETRBase.

• UNETRBase._load_decoder_from_checkpoint(checkpoint), next to _load_encoder_from_checkpoint. It loads all parameters outside the encoder from a path or a state dict (with weights_only=True). It's based on the loading loop in micro_sam's get_unetr.
• A new decoder_checkpoint argument for UNETR and UNETR3D, the same as encoder_checkpoint. The method is called at the end of the subclass __init__, because the decoder does not exist yet in UNETRBase.__init__. Like encoder_checkpoint, it is not stored in init_kwargs.
• get_unetr_model(..., init_decoder=True) then only finds a decoder checkpoint for the given model_type (for example from the micro_sam model registry) and passes it on.

One concern (brought up by Claude) is out_conv. Its shape depends on out_channels, so a pretrained decoder only loads strictly if the new model has the same number of output channels. For example, the micro-sam decoders have out_channels=3, but our synapse-net models use 2. micro_sam's get_unetr solves this with flexible_load_checkpoint: parameters that are missing or have a different shape keep their initialization, with a warning. Should we use the same option here, or load strictly (like the encoder) and only reinitialize out_conv when its shape differs?

[constantinpape]

Regarding your question @stmartineau99 :

Should we use the same option here, or load strictly (like the encoder) and only reinitialize out_conv when its shape differs?

The flexible_load_checkpoint option makes sense here as well. This enables you loading a trained UNETR and use it as initialization for a different segmentation task (with a different number of output channels).

As a general comment: we should align with @anwai98 on Monday to check which functionality should live where.