synapse_net.training.domain_adaptation

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