synapse_net.training.domain_adaptation

  1import os
  2import tempfile
  3from glob import glob
  4from pathlib import Path
  5from typing import Optional, Tuple
  6
  7from sklearn.model_selection import train_test_split
  8
  9import mrcfile
 10import torch
 11import torch_em
 12import torch_em.self_training as self_training
 13from elf.io import open_file
 14
 15from torch_em.transform.invertible_augmentations import MeanTeacherAugmenters, DEFAULT_WEAK_AUGMENTATIONS
 16
 17from .dataloaders import get_supervised_loader, get_unsupervised_loader, _determine_ndim
 18from .models import get_raw_transform, _check_backbone, _get_model
 19from .supervised_training import _derive_key_from_files
 20from ..inference.inference import get_model_path, compute_scale_from_voxel_size, get_available_models
 21from ..inference.util import _restore_legacy_attributes, _Scaler
 22
 23
 24def mean_teacher_adaptation(
 25    name: str,
 26    unsupervised_train_paths: Tuple[str],
 27    unsupervised_val_paths: Tuple[str],
 28    patch_shape: Tuple[int, int, int],
 29    save_root: Optional[str] = None,
 30    source_checkpoint: Optional[str] = None,
 31    supervised_train_paths: Optional[Tuple[str]] = None,
 32    supervised_val_paths: Optional[Tuple[str]] = None,
 33    confidence_threshold: float = 0.9,
 34    raw_key: str = "raw",
 35    raw_key_supervised: str = "raw",
 36    label_key: Optional[str] = None,
 37    batch_size: int = 1,
 38    lr: float = 1e-4,
 39    n_iterations: int = int(1e4),
 40    n_samples_train: Optional[int] = None,
 41    n_samples_val: Optional[int] = None,
 42    train_mask_paths: Optional[Tuple[str]] = None,
 43    val_mask_paths: Optional[Tuple[str]] = None,
 44    sample_mask_key: Optional[str] = None,
 45    unsupervised_sampler: Optional[callable] = None,
 46    supervised_sampler: Optional[callable] = None,
 47    backbone: Optional[str] = None,
 48    model_type: Optional[str] = None,
 49    encoder_checkpoint: Optional[str] = None,
 50    aug_dict: Optional[dict] = None,
 51    separate_backward: bool = False,
 52    check: bool = False,
 53) -> None:
 54    """Run domain adaptation to transfer a network trained on a source domain for a supervised
 55    segmentation task to perform this task on a different target domain.
 56
 57    We support different domain adaptation settings:
 58    - unsupervised domain adaptation: the default mode when 'supervised_train_paths' and
 59     'supervised_val_paths' are not given.
 60    - semi-supervised domain adaptation: domain adaptation on unlabeled and labeled data,
 61      when 'supervised_train_paths' and 'supervised_val_paths' are given.
 62
 63    Args:
 64        name: The name for the checkpoint to be trained.
 65        unsupervised_train_paths: Filepaths to the hdf5 files or similar file formats
 66            for the training data in the target domain.
 67            This training data is used for unsupervised learning, so it does not require labels.
 68        unsupervised_val_paths: Filepaths to the hdf5 files or similar file formats
 69            for the validation data in the target domain.
 70            This validation data is used for unsupervised learning, so it does not require labels.
 71        patch_shape: The patch shape used for a training example.
 72            In order to run 2d training pass a patch shape with a singleton in the z-axis,
 73            e.g. 'patch_shape = [1, 512, 512]'.
 74        save_root: Folder where the checkpoint will be saved.
 75        source_checkpoint: Checkpoint to the initial model trained on the source domain.
 76            This is used to initialize the teacher model.
 77            If the checkpoint is not given, then both student and teacher model are initialized
 78            from scratch. In this case `supervised_train_paths` and `supervised_val_paths` have to
 79            be given in order to provide training data from the source domain.
 80        supervised_train_paths: Filepaths to the hdf5 files for the training data in the source domain.
 81            This training data is optional. If given, it is used for supervised learning and requires labels.
 82        supervised_val_paths: Filepaths to the hdf5 files for the validation data in the source domain.
 83            This validation data is optional. If given, it is used for supervised learning and requires labels.
 84        confidence_threshold: The threshold for filtering data in the unsupervised loss.
 85            The label filtering is done based on the uncertainty of network predictions, and only
 86            the data with higher certainty than this threshold is used for training.
 87        raw_key: The key that holds the raw data inside of the hdf5 or similar files.
 88        raw_key_supervised: The key that holds the raw data inside of the hdf5 files for supervised learning.
 89        label_key: The key that holds the labels inside of the hdf5 files for supervised learning.
 90            This is only required if `supervised_train_paths` and `supervised_val_paths` are given.
 91        batch_size: The batch size for training.
 92        lr: The initial learning rate.
 93        n_iterations: The number of iterations to train for.
 94        n_samples_train: The number of train samples per epoch. By default this will be estimated
 95            based on the patch_shape and size of the volumes used for training.
 96        n_samples_val: The number of val samples per epoch. By default this will be estimated
 97            based on the patch_shape and size of the volumes used for validation.
 98        train_mask_paths: Sample masks used by the unsupervised sampler to accept or reject patches for training.
 99        val_mask_paths: Sample masks used by the unsupervised sampler to accept or reject patches for validation.
100        sample_mask_key: The key to the sample mask dataset inside each file.
101        unsupervised_sampler: Sampler to accept or reject patches for the unsupervised data stream.
102        supervised_sampler: Sampler to accept or reject patches for the supervised data stream.
103            Pass `False` to disable.
104        backbone: The pretrained encoder. Options: "sam", "sam2", "dinov2", "dinov3" (ViT encoders of a UNETR),
105            or "cryosiam" (3D CNN encoder of a CryoZooUNet).
106            The ViT backbones must be set together with `model_type`.
107        model_type: Model type for the selected ViT `backbone` model family, for example "vit_b" or
108            "vit_b_em_organelles" for "sam".
109        encoder_checkpoint: Path to the pretrained encoder weights for `backbone`.
110            By default they are fetched for `backbone` and `model_type`.
111        aug_dict: The augmentations for the teacher and the student, with keys "intensity" and "geometrical".
112            Defaults to Gaussian blur and noise, plus random flips and 90-degree rotations.
113            The noise is scaled with `clip_max` to match the range of `raw_transform`.
114        separate_backward: Whether to backpropagate each loss term separately to reduce peak memory.
115        check: Whether to check the training and validation loaders instead of running training.
116    """
117    assert (supervised_train_paths is None) == (supervised_val_paths is None)
118    _check_backbone(backbone, model_type, patch_shape)
119    is_2d, _ = _determine_ndim(patch_shape)
120
121    if source_checkpoint is None:
122        # training from scratch only makes sense if we have supervised training data
123        # that's why we have the assertion here.
124        assert supervised_train_paths is not None
125        print("Mean teacher training from scratch (AdaMT).")
126        model = _get_model(
127            patch_shape, out_channels=2, backbone=backbone,
128            model_type=model_type, encoder_checkpoint=encoder_checkpoint,
129        )
130        reinit_teacher = True
131    else:
132        print("Mean teacher training initialized from source model:", source_checkpoint)
133        if os.path.isdir(source_checkpoint):
134            model = torch_em.util.load_model(source_checkpoint)
135        else:
136            model = _restore_legacy_attributes(torch.load(source_checkpoint, weights_only=False))
137        reinit_teacher = False
138
139    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
140    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.5, patience=5)
141
142    # self training functionality
143    pseudo_labeler = self_training.DefaultPseudoLabeler(confidence_threshold=confidence_threshold)
144    loss = self_training.SelfTrainingLossWithInvertibleAugmentations()
145    loss_and_metric = self_training.SelfTrainingLossAndMetricWithInvertibleAugmentations()
146
147    raw_transform, clip_max = get_raw_transform(backbone)
148
149    ndim = 2 if is_2d else 3
150    if aug_dict is None:
151        aug_dict = {
152            "intensity": {
153                "RandomGaussianBlur": {"kernel_size": (19, 19), "sigma": (0.1, 3.0)},
154                "RandomGaussianNoise": {"mean": 0.0, "std": 0.1 if clip_max is None else 0.015 * clip_max},
155            },
156            "geometrical": DEFAULT_WEAK_AUGMENTATIONS["geometrical"],
157        }
158    augmenters = MeanTeacherAugmenters(ndim=ndim, clip_max=clip_max, aug_dict=aug_dict)
159
160    unsupervised_train_loader = get_unsupervised_loader(
161        data_paths=unsupervised_train_paths,
162        raw_key=raw_key,
163        patch_shape=patch_shape,
164        batch_size=batch_size,
165        n_samples=n_samples_train,
166        sample_mask_paths=train_mask_paths,
167        sample_mask_key=sample_mask_key,
168        sampler=unsupervised_sampler,
169        raw_transform=raw_transform,
170    )
171    unsupervised_val_loader = get_unsupervised_loader(
172        data_paths=unsupervised_val_paths,
173        raw_key=raw_key,
174        patch_shape=patch_shape,
175        batch_size=batch_size,
176        n_samples=n_samples_val,
177        sample_mask_paths=val_mask_paths,
178        sample_mask_key=sample_mask_key,
179        sampler=unsupervised_sampler,
180        raw_transform=raw_transform,
181    )
182
183    if supervised_train_paths is not None:
184        assert label_key is not None
185        supervised_train_loader = get_supervised_loader(
186            supervised_train_paths, raw_key_supervised, label_key,
187            patch_shape, batch_size, n_samples=n_samples_train,
188            sampler=supervised_sampler, raw_transform=raw_transform,
189        )
190        supervised_val_loader = get_supervised_loader(
191            supervised_val_paths, raw_key_supervised, label_key,
192            patch_shape, batch_size, n_samples=n_samples_val,
193            sampler=supervised_sampler, raw_transform=raw_transform,
194        )
195    else:
196        supervised_train_loader = None
197        supervised_val_loader = None
198
199    if check:
200        from torch_em.util.debug import check_loader
201        check_loader(unsupervised_train_loader, n_samples=4)
202        check_loader(unsupervised_val_loader, n_samples=4)
203        if supervised_train_loader is not None:
204            check_loader(supervised_train_loader, n_samples=4)
205            check_loader(supervised_val_loader, n_samples=4)
206        return
207
208    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
209    trainer = self_training.MeanTeacherTrainerWithInvertibleAugmentations(
210        name=name,
211        model=model,
212        optimizer=optimizer,
213        lr_scheduler=scheduler,
214        pseudo_labeler=pseudo_labeler,
215        unsupervised_loss=loss,
216        unsupervised_loss_and_metric=loss_and_metric,
217        supervised_train_loader=supervised_train_loader,
218        unsupervised_train_loader=unsupervised_train_loader,
219        supervised_val_loader=supervised_val_loader,
220        unsupervised_val_loader=unsupervised_val_loader,
221        supervised_loss=loss,
222        supervised_loss_and_metric=loss_and_metric,
223        logger=self_training.SelfTrainingTensorboardLogger,
224        # On the CPU, mixed precision uses bfloat16. It is very slow without hardware support,
225        # and the kornia augmentations do not accept it.
226        mixed_precision=device.type == "cuda",
227        log_image_interval=100,
228        compile_model=False,
229        device=device,
230        reinit_teacher=reinit_teacher,
231        save_root=save_root,
232        augmenter=augmenters,
233        separate_backward=separate_backward,
234    )
235    trainer.fit(n_iterations)
236
237
238# TODO patch shapes for other models
239PATCH_SHAPES = {
240    "vesicles_3d": [48, 256, 256],
241}
242"""@private
243"""
244
245
246def _get_paths(input_folder, pattern, resize_training_data, model_name, tmp_dir, val_fraction):
247    files = sorted(glob(os.path.join(input_folder, "**", pattern), recursive=True))
248    if len(files) == 0:
249        raise ValueError(f"Could not load any files from {input_folder} with pattern {pattern}")
250
251    # Heuristic: if we have less then 4 files then we crop a part of the volumes for validation.
252    # And resave the volumes.
253    resave_val_crops = len(files) < 4
254
255    # We only resave the data if we resave val crops or resize the training data
256    resave_data = resave_val_crops or resize_training_data
257    if not resave_data:
258        train_paths, val_paths = train_test_split(files, test_size=val_fraction)
259        return train_paths, val_paths
260
261    train_paths, val_paths = [], []
262    for file_path in files:
263        file_name = os.path.basename(file_path)
264        data = open_file(file_path, mode="r")["data"][:]
265
266        if resize_training_data:
267            with mrcfile.open(file_path) as f:
268                voxel_size = f.voxel_size
269            voxel_size = {ax: vox_size / 10.0 for ax, vox_size in zip("xyz", voxel_size.item())}
270            scale = compute_scale_from_voxel_size(voxel_size, model_name)
271            scaler = _Scaler(scale, verbose=False)
272            data = scaler.scale_input(data)
273
274        if resave_val_crops:
275            n_slices = data.shape[0]
276            val_slice = int((1.0 - val_fraction) * n_slices)
277            train_data, val_data = data[:val_slice], data[val_slice:]
278
279            train_path = os.path.join(tmp_dir, Path(file_name).with_suffix(".h5")).replace(".h5", "_train.h5")
280            with open_file(train_path, mode="w") as f:
281                f.create_dataset("data", data=train_data, compression="lzf")
282            train_paths.append(train_path)
283
284            val_path = os.path.join(tmp_dir, Path(file_name).with_suffix(".h5")).replace(".h5", "_val.h5")
285            with open_file(val_path, mode="w") as f:
286                f.create_dataset("data", data=val_data, compression="lzf")
287            val_paths.append(val_path)
288
289        else:
290            output_path = os.path.join(tmp_dir, Path(file_name).with_suffix(".h5"))
291            with open_file(output_path, mode="w") as f:
292                f.create_dataset("data", data=data, compression="lzf")
293            train_paths.append(output_path)
294
295    if not resave_val_crops:
296        train_paths, val_paths = train_test_split(train_paths, test_size=val_fraction)
297
298    return train_paths, val_paths
299
300
301def _parse_patch_shape(patch_shape, model_name):
302    if patch_shape is None:
303        patch_shape = PATCH_SHAPES[model_name]
304    return patch_shape
305
306
307def _resolve_source_checkpoint(source_model, source_checkpoint):
308    if source_checkpoint is None:
309        return get_model_path(source_model)
310    if not os.path.exists(source_checkpoint):
311        raise ValueError(f"The source checkpoint does not exist: {source_checkpoint}")
312    return source_checkpoint
313
314
315def main():
316    """@private
317    """
318    import argparse
319
320    parser = argparse.ArgumentParser(
321        description="Adapt a model to data from a different domain using unsupervised domain adaptation.\n\n"
322        "You can use this function to adapt the SynapseNet model for vesicle segmentation like this:\n"
323        "synapse_net.run_domain_adaptation -n adapted_model -i /path/to/data --file_pattern *.mrc --source_model vesicles_3d\n"  # noqa
324        "The trained model will be saved in the folder 'checkpoints/adapted_model' (or whichever name you pass to the '-n' argument)."  # noqa
325        "You can then use this model for segmentation with the SynapseNet GUI or CLI. "
326        "Check out the information below for details on the arguments of this function.",
327        formatter_class=argparse.RawTextHelpFormatter
328    )
329    parser.add_argument("--name", "-n", required=True, help="The name of the model to be trained. ")
330    parser.add_argument("--input_folder", "-i", required=True, help="The folder with the training data.")
331    parser.add_argument("--file_pattern", default="*",
332                        help="The pattern for selecting files for training. For example '*.mrc' to select mrc files.")
333    parser.add_argument("--key", help="The internal file path for the training data. Will be derived from the file extension by default.")  # noqa
334    available_models = get_available_models()
335    parser.add_argument(
336        "--source_model",
337        default="vesicles_3d",
338        help="The source model sets the default checkpoint and the metadata for patch size and data resizing. "
339        "By default the model 'vesicles_3d' for vesicle segmentation in volumetric data is used.\n"
340        "The --source_checkpoint argument overrides the default checkpoint.\n"
341        f"The following source models are available: {available_models}"
342    )
343    parser.add_argument(
344        "--source_checkpoint",
345        help="Path to a custom source checkpoint directory or serialized PyTorch model. "
346        "If not given, the command downloads the checkpoint for --source_model."
347    )
348    parser.add_argument(
349        "--resize_training_data", action="store_true",
350        help="Whether to resize the training data to fit the voxel size of the source model's trainign data."
351    )
352    parser.add_argument("--n_iterations", type=int, default=int(1e4), help="The number of iterations for training.")
353    parser.add_argument(
354        "--patch_shape", nargs=3, type=int,
355        help="The patch shape for training. By default the patch shape the source model was trained with is used."
356    )
357
358    # More optional argument:
359    parser.add_argument("--batch_size", type=int, default=1, help="The batch size for training.")
360    parser.add_argument("--n_samples_train", type=int, help="The number of samples per epoch for training. If not given will be derived from the data size.")  # noqa
361    parser.add_argument("--n_samples_val", type=int, help="The number of samples per epoch for validation. If not given will be derived from the data size.")  # noqa
362    parser.add_argument("--val_fraction", type=float, default=0.15, help="The fraction of the data to use for validation. This has no effect if 'val_folder' and 'val_label_folder' were passed.")  # noqa
363    parser.add_argument("--check", action="store_true", help="Visualize samples from the data loaders to ensure correct data instead of running training.")  # noqa
364    parser.add_argument("--save_root", help="Root path for saving the checkpoint and log dir.")
365
366    args = parser.parse_args()
367
368    try:
369        source_checkpoint = _resolve_source_checkpoint(args.source_model, args.source_checkpoint)
370    except ValueError as error:
371        parser.error(str(error))
372    patch_shape = _parse_patch_shape(args.patch_shape, args.source_model)
373    with tempfile.TemporaryDirectory() as tmp_dir:
374        unsupervised_train_paths, unsupervised_val_paths = _get_paths(
375            args.input_folder, args.file_pattern, args.resize_training_data,
376            args.source_model, tmp_dir, args.val_fraction,
377        )
378        unsupervised_train_paths, raw_key = _derive_key_from_files(unsupervised_train_paths, args.key)
379
380        mean_teacher_adaptation(
381            name=args.name,
382            unsupervised_train_paths=unsupervised_train_paths,
383            unsupervised_val_paths=unsupervised_val_paths,
384            patch_shape=patch_shape,
385            source_checkpoint=source_checkpoint,
386            raw_key=raw_key,
387            n_iterations=args.n_iterations,
388            batch_size=args.batch_size,
389            n_samples_train=args.n_samples_train,
390            n_samples_val=args.n_samples_val,
391            check=args.check,
392            save_root=args.save_root,
393        )