synapse_net.training.semisupervised_training

  1import os
  2from typing import Optional, Tuple
  3
  4from .domain_adaptation import mean_teacher_adaptation
  5from .supervised_training import supervised_training
  6
  7
  8def semisupervised_training(
  9    name: str,
 10    unsupervised_train_paths: Tuple[str],
 11    unsupervised_val_paths: Tuple[str],
 12    patch_shape: Tuple[int, int, int],
 13    save_root: str,
 14    supervised_train_paths: Tuple[str],
 15    supervised_val_paths: Tuple[str],
 16    label_key: str,
 17    source_checkpoint=None,
 18    confidence_threshold: float = 0.9,
 19    raw_key: str = "raw",
 20    batch_size: int = 1,
 21    lr: float = 1e-4,
 22    n_iterations: int = int(1e5),
 23    teacher_warmup_iterations: int = int(1e4),
 24    n_samples_train: Optional[int] = None,
 25    n_samples_val: Optional[int] = None,
 26    train_mask_paths: Optional[Tuple[str]] = None,
 27    val_mask_paths: Optional[Tuple[str]] = None,
 28    sample_mask_key: Optional[str] = None,
 29    unsupervised_sampler: Optional[callable] = None,
 30    supervised_sampler: Optional[callable] = None,
 31    backbone: Optional[str] = None,
 32    model_type: Optional[str] = None,
 33    encoder_checkpoint: Optional[str] = None,
 34    aug_dict: Optional[dict] = None,
 35    separate_backward: bool = False,
 36    check: bool = False,
 37):
 38    """Run semisupervised segmentation training.
 39
 40    This proceeds in two phases:
 41
 42        1. If no `source_checkpoint` is given, run supervised training to
 43            warmup the teacher.
 44        2. Run semisupervised training using mean teacher setup with invertible
 45            augmentations, using the warmup checkpoint to initialize the teacher.
 46
 47    Args:
 48        name: The name for the checkpoint to be trained. The warmup checkpoint is saved
 49            under the name "{name}-warmup".
 50        unsupervised_train_paths: Filepaths to the hdf5 files for the unsupervised
 51            training data. This data does not require labels.
 52        unsupervised_val_paths: Filepaths to the hdf5 files for the unsupervised
 53            validation data. This data does not require labels.
 54        patch_shape: The patch shape used for a training example.
 55            In order to run 2d training pass a patch shape with a singleton in the z-axis,
 56            e.g. 'patch_shape = [1, 512, 512]'.
 57        save_root: Folder where the checkpoints will be saved.
 58        supervised_train_paths: Filepaths to the hdf5 files for the supervised training
 59            data, requires labels. Used for both the teacher warmup and the semi-supervised loss.
 60        supervised_val_paths: Filepaths to the hdf5 files for the supervised validation data,
 61            requires labels.
 62        label_key: The key that holds the labels inside of the hdf5 files.
 63        source_checkpoint: Warmup checkpoint used to initialize the teacher model. If not provided,
 64            run supervised training for `teacher_warmup_iterations`.
 65        confidence_threshold: The threshold for filtering data in the unsupervised loss.
 66            The label filtering is done based on the uncertainty of network predictions, and only
 67            the data with higher certainty than this threshold is used for training.
 68        raw_key: The key that holds the raw data inside of the hdf5 files.
 69        batch_size: The batch size for training.
 70        lr: The initial learning rate.
 71        n_iterations: The number of mean-teacher training iterations.
 72            Set to 0 to stop after the teacher warmup.
 73        teacher_warmup_iterations: The number of iterations for the supervised teacher warmup,
 74            only used if no warmup checkpoint exists yet and no `source_checkpoint` is given.
 75        n_samples_train: The number of train samples per epoch. By default this will be estimated
 76            based on the patch_shape and size of the volumes used for training.
 77        n_samples_val: The number of val samples per epoch. By default this will be estimated
 78            based on the patch_shape and size of the volumes used for validation.
 79        train_mask_paths: Sample masks used by the unsupervised sampler to accept or reject patches for training.
 80        val_mask_paths: Sample masks used by the unsupervised sampler to accept or reject patches for validation.
 81        sample_mask_key: The key to the sample mask dataset inside each file.
 82        unsupervised_sampler: Sampler to accept or reject patches for the unsupervised data stream.
 83        supervised_sampler: Sampler to accept or reject patches for the supervised data stream.
 84        backbone: The pretrained encoder. Options: "sam", "sam2", "dinov2", "dinov3" (ViT encoders of a UNETR),
 85            or "cryosiam" (3D CNN encoder of a CryoZooUNet).
 86            The ViT backbones must be set together with `model_type`.
 87        model_type: Model type for the selected ViT `backbone` model family, for example "vit_b" or
 88            "vit_b_em_organelles" for "sam".
 89        encoder_checkpoint: Path to the pretrained encoder weights for `backbone`.
 90            By default they are fetched for `backbone` and `model_type`.
 91        aug_dict: The augmentations for the teacher and the student, with keys "intensity" and "geometrical".
 92            Defaults to Gaussian blur and noise, plus random flips and 90-degree rotations.
 93        separate_backward: Whether to backpropagate each loss term separately to reduce peak memory.
 94        check: Whether to check the training and validation loaders instead of running training.
 95    """
 96    if check:
 97        mean_teacher_adaptation(
 98            name=name,
 99            unsupervised_train_paths=unsupervised_train_paths,
100            unsupervised_val_paths=unsupervised_val_paths,
101            patch_shape=patch_shape,
102            save_root=save_root,
103            source_checkpoint=source_checkpoint,
104            supervised_train_paths=supervised_train_paths,
105            supervised_val_paths=supervised_val_paths,
106            confidence_threshold=confidence_threshold,
107            raw_key=raw_key,
108            raw_key_supervised=raw_key,
109            label_key=label_key,
110            batch_size=batch_size,
111            lr=lr,
112            n_iterations=n_iterations,
113            n_samples_train=n_samples_train,
114            n_samples_val=n_samples_val,
115            train_mask_paths=train_mask_paths,
116            val_mask_paths=val_mask_paths,
117            sample_mask_key=sample_mask_key,
118            unsupervised_sampler=unsupervised_sampler,
119            supervised_sampler=supervised_sampler,
120            backbone=backbone,
121            model_type=model_type,
122            encoder_checkpoint=encoder_checkpoint,
123            aug_dict=aug_dict,
124            separate_backward=separate_backward,
125            check=True,
126        )
127        return
128
129    warmup_name = f"{name}-warmup"
130    if source_checkpoint is None:
131        warmup_checkpoint = os.path.join(save_root, "checkpoints", warmup_name, "best.pt")
132
133        if not os.path.exists(warmup_checkpoint):
134            print(
135                "No warmup checkpoint was found, initiating supervised warmup for teacher model "
136                f"with {teacher_warmup_iterations} iterations."
137            )
138
139            supervised_training(
140                name=warmup_name,
141                train_paths=supervised_train_paths,
142                val_paths=supervised_val_paths,
143                label_key=label_key,
144                patch_shape=patch_shape,
145                save_root=save_root,
146                raw_key=raw_key,
147                batch_size=batch_size,
148                lr=lr,
149                n_iterations=teacher_warmup_iterations,
150                sampler=supervised_sampler,
151                n_samples_train=n_samples_train,
152                n_samples_val=n_samples_val,
153                backbone=backbone,
154                model_type=model_type,
155                encoder_checkpoint=encoder_checkpoint,
156                check=False,
157            )
158        source_checkpoint = os.path.dirname(warmup_checkpoint)
159
160    if n_iterations == 0:
161        return
162
163    mean_teacher_adaptation(
164        name=name,
165        unsupervised_train_paths=unsupervised_train_paths,
166        unsupervised_val_paths=unsupervised_val_paths,
167        patch_shape=patch_shape,
168        save_root=save_root,
169        source_checkpoint=source_checkpoint,
170        supervised_train_paths=supervised_train_paths,
171        supervised_val_paths=supervised_val_paths,
172        confidence_threshold=confidence_threshold,
173        raw_key=raw_key,
174        raw_key_supervised=raw_key,
175        label_key=label_key,
176        batch_size=batch_size,
177        lr=lr,
178        n_iterations=n_iterations,
179        n_samples_train=n_samples_train,
180        n_samples_val=n_samples_val,
181        train_mask_paths=train_mask_paths,
182        val_mask_paths=val_mask_paths,
183        sample_mask_key=sample_mask_key,
184        unsupervised_sampler=unsupervised_sampler,
185        supervised_sampler=supervised_sampler,
186        backbone=backbone,
187        model_type=model_type,
188        encoder_checkpoint=encoder_checkpoint,
189        aug_dict=aug_dict,
190        separate_backward=separate_backward,
191        check=False,
192    )