SSDA changes (#168)

* update domain_adaptation.py to use

* formatting changes +  arg added to supervised_learning.py

* respond to PR revssda iew comments

---------

Co-authored-by: GOESTERN-1107078 <u25914@glogin10.usr.hpc.gwdg.de>

Author: stmartineau99

synapse_net/training/domain_adaptation.py

@@ -18,6 +18,14 @@
 from ..inference.inference import get_model_path, compute_scale_from_voxel_size, get_available_models
 from ..inference.util import _Scaler
 
+# configure weak augmentations
+from torch_em.transform.invertible_augmentations import DEFAULT_WEAK_AUGMENTATIONS
+
+# TODO - test settings - fixed kernel size for `RandomGaussianBlur`, fixed std for `RandomGaussianBlur`
+DEFAULT_WEAK_AUGMENTATIONS["intensity"] = {
+    "RandomGaussianBlur": {"kernel_size": (19, 19), "sigma": (0.1, 3.0)},
+    "RandomGaussianNoise": {"mean": (0.0), "std": (0.1)},
+}    
 
 def mean_teacher_adaptation(
     name: str,
@@ -40,7 +48,8 @@ def mean_teacher_adaptation(
     train_mask_paths: Optional[Tuple[str]] = None,
     val_mask_paths: Optional[Tuple[str]] = None,
     sample_mask_key: Optional[str] = None,
-    patch_sampler: Optional[callable] = None,
+    unsupervised_sampler: Optional[callable] = None,
+    supervised_sampler: Optional[callable] = None,
     check: bool = False,
 ) -> None:
     """Run domain adaptation to transfer a network trained on a source domain for a supervised
@@ -86,10 +95,12 @@ def mean_teacher_adaptation(
             based on the patch_shape and size of the volumes used for training.
         n_samples_val: The number of val samples per epoch. By default this will be estimated
             based on the patch_shape and size of the volumes used for validation.
-        train_mask_paths: Sample masks used by the patch sampler to accept or reject patches for training.
-        val_mask_paths: Sample masks used by the patch sampler to accept or reject patches for validation.
+        train_mask_paths: Sample masks used by the unsupervised sampler to accept or reject patches for training.
+        val_mask_paths: Sample masks used by the unsupervised sampler to accept or reject patches for validation.
         sample_mask_key: The key to the sample mask dataset inside each file.
-        patch_sampler: Accept or reject patches based on a condition.
+        unsupervised_sampler: Sampler to accept or reject patches for the unsupervised data stream.
+        supervised_sampler: Sampler to accept or reject patches for the supervised data stream.
+            Pass `False` to disable.
         check: Whether to check the training and validation loaders instead of running training.
     """  # noqa
     assert (supervised_train_paths is None) == (supervised_val_paths is None)
@@ -118,8 +129,11 @@ def mean_teacher_adaptation(
 
     # self training functionality
     pseudo_labeler = self_training.DefaultPseudoLabeler(confidence_threshold=confidence_threshold)
-    loss = self_training.DefaultSelfTrainingLoss()
-    loss_and_metric = self_training.DefaultSelfTrainingLossAndMetric()
+    loss = self_training.SelfTrainingLossWithInvertibleAugmentations()
+    loss_and_metric = self_training.SelfTrainingLossAndMetricWithInvertibleAugmentations()
+
+    ndim = 2 if is_2d else 3
+    augmenters = torch_em.transform.invertible_augmentations.MeanTeacherAugmenters(ndim=ndim)
 
     unsupervised_train_loader = get_unsupervised_loader(
         data_paths=unsupervised_train_paths,
@@ -129,7 +143,7 @@ def mean_teacher_adaptation(
         n_samples=n_samples_train,
         sample_mask_paths=train_mask_paths,
         sample_mask_key=sample_mask_key,
-        sampler=patch_sampler,
+        sampler=unsupervised_sampler,
     )
     unsupervised_val_loader = get_unsupervised_loader(
         data_paths=unsupervised_val_paths,
@@ -139,18 +153,20 @@ def mean_teacher_adaptation(
         n_samples=n_samples_val,
         sample_mask_paths=val_mask_paths,
         sample_mask_key=sample_mask_key,
-        sampler=patch_sampler,
+        sampler=unsupervised_sampler,
     )
 
     if supervised_train_paths is not None:
         assert label_key is not None
         supervised_train_loader = get_supervised_loader(
             supervised_train_paths, raw_key_supervised, label_key,
             patch_shape, batch_size, n_samples=n_samples_train,
+            sampler=supervised_sampler,
         )
         supervised_val_loader = get_supervised_loader(
             supervised_val_paths, raw_key_supervised, label_key,
             patch_shape, batch_size, n_samples=n_samples_val,
+            sampler=supervised_sampler,
         )
     else:
         supervised_train_loader = None
@@ -166,7 +182,7 @@ def mean_teacher_adaptation(
         return
 
     device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-    trainer = self_training.MeanTeacherTrainer(
+    trainer = self_training.MeanTeacherTrainerWithInvertibleAugmentations(
         name=name,
         model=model,
         optimizer=optimizer,
@@ -187,6 +203,7 @@ def mean_teacher_adaptation(
         device=device,
         reinit_teacher=reinit_teacher,
         save_root=save_root,
+        augmenter=augmenters,
     )
     trainer.fit(n_iterations)
 

synapse_net/training/semisupervised_training.py

@@ -57,7 +57,7 @@ def get_unsupervised_loader(
         sample_mask_key: The key to the sample mask dataset inside each file.
         bg_mask_paths: The filepaths to the background masks for each tomogram.
         bg_mask_key: The key to the background mask dataset inside each file.
-        sampler: Accept or reject patches based on a condition.
+        sampler: Optional sampler to accept or reject patches for training. 
         exclude_top_and_bottom: Whether to exclude the five top and bottom slices to
             avoid artifacts at the border of tomograms.
 
@@ -79,7 +79,7 @@ def get_unsupervised_loader(
     _, ndim = _determine_ndim(patch_shape)
     raw_transform = torch_em.transform.get_raw_transform()
     transform = torch_em.transform.get_augmentations(ndim=ndim)
-    augmentations = (weak_augmentations(), weak_augmentations())
+    # augmentations = (weak_augmentations(), weak_augmentations())
 
     if n_samples is None:
         n_samples_per_ds = None
@@ -97,7 +97,7 @@ def get_unsupervised_loader(
             n_samples=n_samples_per_ds,
             sampler=sampler,
             ndim=ndim,
-            augmentations=augmentations,
+            augmentations=None,
             sample_mask_path=sample_mask_paths[i] if sample_mask_paths is not None else None,
             sample_mask_key=sample_mask_key,
             bg_mask_path=bg_mask_paths[i] if bg_mask_paths is not None else None,

synapse_net/training/supervised_training.py

@@ -1,6 +1,6 @@
 import os
 from glob import glob
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Union
 
 import torch
 import torch_em
@@ -97,7 +97,7 @@ def get_supervised_loader(
     add_boundary_transform: bool = True,
     label_dtype=torch.float32,
     rois: Optional[Tuple[Tuple[slice]]] = None,
-    sampler: Optional[callable] = None,
+    sampler: Optional[Union[callable, bool]] = None,
     ignore_label: Optional[int] = None,
     label_transform: Optional[callable] = None,
     label_paths: Optional[Tuple[str]] = None,
@@ -118,8 +118,8 @@ def get_supervised_loader(
         add_boundary_transform: Whether to add a boundary channel to the training data.
         label_dtype: The datatype of the labels returned by the dataloader.
         rois: Optional region of interests for training.
-        sampler: Optional sampler for selecting blocks for training.
-            By default a minimum instance sampler will be used.
+        sampler: Optional sampler to accept or reject patches for training. 
+            By default a minimum instance sampler will be used, pass `False` to disable.
         ignore_label: Ignore label in the ground-truth. The areas marked by this label will be
             ignored in the loss computation. By default this option is not used.
         label_transform: Label transform that is applied to the segmentation to compute the targets.
@@ -162,6 +162,8 @@ def get_supervised_loader(
 
     if sampler is None:
         sampler = torch_em.data.sampler.MinInstanceSampler(min_num_instances=4)
+    elif sampler is False:
+        sampler = None
 
     if label_paths is None:
         label_paths = data_paths
@@ -194,7 +196,7 @@ def supervised_training(
     val_label_paths: Optional[Tuple[str]] = None,
     train_rois: Optional[Tuple[Tuple[slice]]] = None,
     val_rois: Optional[Tuple[Tuple[slice]]] = None,
-    sampler: Optional[callable] = None,
+    sampler: Optional[Union[callable, bool]] = None,
     n_samples_train: Optional[int] = None,
     n_samples_val: Optional[int] = None,
     check: bool = False,
@@ -205,6 +207,7 @@ def supervised_training(
     out_channels: int = 2,
     mask_channel: bool = False,
     checkpoint_path: Optional[str] = None,
+    save_every_kth_epoch: Optional[int] = None,
     **loader_kwargs,
 ):
     """Run supervised segmentation training.
@@ -233,7 +236,7 @@ def supervised_training(
             If not given, the labels are expected to be part of `val_paths`.
         train_rois: Optional region of interests for training.
         val_rois: Optional region of interests for validation.
-        sampler: Optional sampler for selecting blocks for training.
+        sampler: Optional sampler for selecting patches for training.
             By default a minimum instance sampler will be used.
         n_samples_train: The number of train samples per epoch. By default this will be estimated
             based on the patch_shape and size of the volumes used for training.
@@ -249,6 +252,8 @@ def supervised_training(
         mask_channel: Whether the last channels in the labels should be used for masking the loss.
             This can be used to implement more complex masking operations and is not compatible with `ignore_label`.
         checkpoint_path: Path to the directory where 'best.pt' resides; continue training this model.
+        save_every_kth_epoch: Save checkpoints after every kth epoch in a separate file.
+            The corresponding checkpoints will be saved with the naming scheme 'epoch-{epoch}.pt'.
         loader_kwargs: Additional keyword arguments for the dataloader.
     """
     train_loader = get_supervised_loader(train_paths, raw_key, label_key, patch_shape, batch_size,
@@ -317,7 +322,7 @@ def supervised_training(
         loss=loss,
         metric=metric,
     )
-    trainer.fit(n_iterations)
+    trainer.fit(n_iterations, save_every_kth_epoch=save_every_kth_epoch)
 
 
 def _derive_key_from_files(files, key):