Add examples for network training (#61)

Add examples for network training

Author: constantinpape

examples/domain_adaptation.py

@@ -0,0 +1,81 @@
+"""This script contains an example for using domain adptation to
+transfer a trained model for vesicle segmentation to a new dataset from a different data distribution,
+e.g. data from regular transmission electron microscopy (2D) instead of electron tomography or data from
+a different electron tomogram with different specimen and sample preparation.
+You don't need any annotations in the new domain to run this script.
+
+You can download example data for this script from:
+- Adaptation to 2d TEM data: TODO zenodo link
+- Adaptation to different tomography data: TODO zenodo link
+"""
+
+import os
+from glob import glob
+
+from sklearn.model_selection import train_test_split
+from synaptic_reconstruction.training import mean_teacher_adaptation
+from synaptic_reconstruction.tools.util import get_model_path
+
+
+def main():
+    # Choose whether to adapt the model to 2D or to 3D data.
+    train_2d_model = True
+
+    # TODO adjust to zenodo downloads
+    # These are the data folders for the example data downloaded from zenodo.
+    # Update these paths to apply the script to your own data.
+    # Check out the example data to see the data format for training.
+    data_root_folder_2d = "./data/2d_tem/train_unlabeled"
+    data_root_folder_3d = "./data/..."
+
+    # Choose the correct data folder depending on 2d/3d training.
+    data_root_folder = data_root_folder_2d if train_2d_model else data_root_folder_3d
+
+    # Get all files with ending .h5 in the training folder.
+    files = sorted(glob(os.path.join(data_root_folder, "**", "*.h5"), recursive=True))
+
+    # Crate a train / val split.
+    train_ratio = 0.85
+    train_paths, val_paths = train_test_split(files, test_size=1 - train_ratio, shuffle=True, random_state=42)
+
+    # Choose settings for the 2d or 3d domain adaptation.
+    if train_2d_model:
+        # This is the name of the checkpoint of the adapted model.
+        # For the name here the checkpoint will be stored in './checkpoints/example-2d-adapted-model'
+        model_name = "example-2d-adapted-model"
+        # The training patch size.
+        patch_shape = (256, 256)
+        # The batch size for training. You can increase this if you have enough VRAM.
+        batch_size = 4
+        # Get the checkpoint of the pretrained model for 2d vesicle segmentation.
+        source_checkpoint = get_model_path(model_type="vesicles_2d")
+    else:
+        # This is the name of the checkpoint of the adapted model.
+        # For the name here the checkpoint will be stored in './checkpoints/example-3d-adapted-model'
+        model_name = "example-3d-adapted-model"
+        # The training patch size.
+        patch_shape = (48, 256, 256)
+        # The batch size for training. You can increase this if you have enough VRAM.
+        batch_size = 1
+        # Get the checkpoint of the pretrained model for d vesicle segmentation.
+        source_checkpoint = get_model_path(model_type="vesicles_3d")
+
+    # We set the number of training iterations to 25,000.
+    n_iterations = int(2.5e4)
+
+    # This function runs the domain adaptation. Check out its documentation for
+    # advanced settings to update the training procedure.
+    mean_teacher_adaptation(
+        name=model_name,
+        unsupervised_train_paths=train_paths,
+        unsupervised_val_paths=val_paths,
+        source_checkpoint=source_checkpoint,
+        patch_shape=patch_shape,
+        batch_size=batch_size,
+        n_iterations=n_iterations,
+        confidence_threshold=0.75,
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/network_training.py

@@ -0,0 +1,77 @@
+"""This script contains an example for how to train a network for
+a segmentation task with SynapseNet. This script covers the case of
+supervised training, i.e. your data needs to contain annotations for
+the structures you want to segment. If you want to use domain adaptation
+to adapt an already trained network to your data without the need for
+additional annotations then check out `domain_adaptation.py`.
+
+You can download example data for this script from:
+TODO zenodo link to Single-Ax / Chemical Fix data.
+"""
+import os
+from glob import glob
+
+from sklearn.model_selection import train_test_split
+from synaptic_reconstruction.training import supervised_training
+
+
+def main():
+    # This is the folder that contains your training data.
+    # The example was designed so that it runs for the sample data downloaded to './data'.
+    # If you want to train on your own data than change this filepath accordingly.
+    # TODO update to match zenodo download
+    data_root_folder = "./data/vesicles/train"
+
+    # The training data should be saved as .h5 files, with:
+    # an internal dataset called 'raw' that contains the image data
+    # and another dataset that contains the training annotations.
+    label_key = "labels/vesicles"
+
+    # Get all files with the ending .h5 in the training folder.
+    files = sorted(glob(os.path.join(data_root_folder, "**", "*.h5"), recursive=True))
+
+    # Crate a train / val split.
+    train_ratio = 0.85
+    train_paths, val_paths = train_test_split(files, test_size=1 - train_ratio, shuffle=True, random_state=42)
+
+    # We can either train a 2d or a 3d model. Whether a 2d or a 3d model is trained is derived from the patch shape.
+    # If your training data for 2d is stored as images (i.e. 2d data) them choose a  patch shape of form Y x X,
+    # e.g. (384, 384). If your data is stored in 3d, but you want to train a 2d model on it, choose a patch shape
+    # of the form 1 x Y x X, e.g. (1, 384, 384).
+    # If you want to train a 3d model then choose a patch shape of form Z x Y x X, e.g. (48, 256, 256).
+    train_2d_model = True
+    if train_2d_model:
+        batch_size = 2  # You can increase the batch size if you have enough VRAM.
+        # The model name determines the name of the checkpoint. E.g., for the name here the checkpoint will
+        # be saved at: 'checkpoints/example-2d-vesicle-model/'.
+        model_name = "example-2d-vesicle-model"
+        # The patch shape for training. See futher explanations above.
+        patch_shape = (1, 384, 384)
+    else:
+        batch_size = 1  # You can increase the batch size if you have enough VRAM.
+        # See the explanations for model_name and patch_shape above.
+        model_name = "example-3d-vesicle-model"
+        patch_shape = (48, 256, 256)
+
+    # If check_loader is set to True the training samples will be visualized via napari
+    # instead of starting a training. This is useful to validate that the training data
+    # is read correctly.
+    check_loader = False
+
+    # This function runs the training. Check out its documentation for
+    # advanced settings to update the training procedure.
+    supervised_training(
+        name=model_name,
+        train_paths=train_paths,
+        val_paths=val_paths,
+        label_key=label_key,
+        patch_shape=patch_shape,
+        batch_size=batch_size,
+        n_samples_train=None,
+        n_samples_val=25,
+        check=check_loader,
+    )
+
+
+if __name__ == "__main__":
+    main()

setup.py

@@ -13,7 +13,7 @@
     license="MIT",
     entry_points={
         "console_scripts": [
-            "synapse_net.run_segmentation = synaptic_reconstruction.tools.cli:segmentation_cli"
+            "synapse_net.run_segmentation = synaptic_reconstruction.tools.cli:segmentation_cli",
         ],
         "napari.manifest": [
             "synaptic_reconstruction = synaptic_reconstruction:napari.yaml",

synaptic_reconstruction/tools/util.py

@@ -10,6 +10,20 @@
 from ..inference.mitochondria import segment_mitochondria
 
 
+def get_model_path(model_type: str) -> str:
+    """Get the local path to a given model.
+
+    Args:
+        The model type.
+
+    Returns:
+        The local path to the model.
+    """
+    model_registry = get_model_registry()
+    model_path = model_registry.fetch(model_type)
+    return model_path
+
+
 def get_model(model_type: str, device: Optional[Union[str, torch.device]] = None) -> torch.nn.Module:
     """Get the model for the given segmentation type.
 
@@ -22,8 +36,7 @@ def get_model(model_type: str, device: Optional[Union[str, torch.device]] = None
         The model.
     """
     device = get_device(device)
-    model_registry = get_model_registry()
-    model_path = model_registry.fetch(model_type)
+    model_path = get_model_path(model_type)
     warnings.filterwarnings(
         "ignore",
         message="You are using `torch.load` with `weights_only=False`",

synaptic_reconstruction/training/domain_adaptation.py

@@ -1,19 +1,20 @@
+import os
 from typing import Optional, Tuple
 
 import torch
 import torch_em
 import torch_em.self_training as self_training
 
 from .semisupervised_training import get_unsupervised_loader
-from .supervised_training import get_2d_model, get_3d_model, get_supervised_loader
+from .supervised_training import get_2d_model, get_3d_model, get_supervised_loader, determine_ndim
 
 
 def mean_teacher_adaptation(
     name: str,
     unsupervised_train_paths: Tuple[str],
     unsupervised_val_paths: Tuple[str],
     patch_shape: Tuple[int, int, int],
-    save_root: str,
+    save_root: Optional[str] = None,
     source_checkpoint: Optional[str] = None,
     supervised_train_paths: Optional[Tuple[str]] = None,
     supervised_val_paths: Optional[Tuple[str]] = None,
@@ -70,22 +71,24 @@ def mean_teacher_adaptation(
             based on the patch_shape and size of the volumes used for validation.
     """
     assert (supervised_train_paths is None) == (supervised_val_paths is None)
+    is_2d, _ = determine_ndim(patch_shape)
 
     if source_checkpoint is None:
         # training from scratch only makes sense if we have supervised training data
         # that's why we have the assertion here.
         assert supervised_train_paths is not None
         print("Mean teacher training from scratch (AdaMT)")
-        # TODO determine 2d vs 3d
-        is_2d = False
         if is_2d:
             model = get_2d_model(out_channels=2)
         else:
             model = get_3d_model(out_channels=2)
         reinit_teacher = True
     else:
         print("Mean teacehr training initialized from source model:", source_checkpoint)
-        model = torch_em.util.load_model(source_checkpoint)
+        if os.path.isdir(source_checkpoint):
+            model = torch_em.util.load_model(source_checkpoint)
+        else:
+            model = torch.load(source_checkpoint)
         reinit_teacher = False
 
     optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

synaptic_reconstruction/training/semisupervised_training.py

@@ -6,7 +6,7 @@
 import torch_em.self_training as self_training
 from torchvision import transforms
 
-from .supervised_training import get_2d_model, get_3d_model, get_supervised_loader
+from .supervised_training import get_2d_model, get_3d_model, get_supervised_loader, determine_ndim
 
 
 def weak_augmentations(p: float = 0.75) -> callable:
@@ -61,14 +61,7 @@ def get_unsupervised_loader(
     else:
         roi = None
 
-    if len(patch_shape) == 2:
-        ndim = 2
-    else:
-        assert len(patch_shape) == 3
-        z, y, x = patch_shape
-        ndim = 2 if z == 1 else 3
-    print("ndim is: ", ndim)
-
+    _, ndim = determine_ndim(patch_shape)
     raw_transform = torch_em.transform.get_raw_transform()
     transform = torch_em.transform.get_augmentations(ndim=ndim)
 

synaptic_reconstruction/training/supervised_training.py

@@ -67,6 +67,18 @@ def adjust_patch_shape(data_shape, patch_shape):
     return patch_shape  # Return the original patch_shape for 3D data
 
 
+def determine_ndim(patch_shape):
+    # Check for 2D or 3D training
+    try:
+        z, y, x = patch_shape
+    except ValueError:
+        y, x = patch_shape
+        z = 1
+    is_2d = z == 1
+    ndim = 2 if is_2d else 3
+    return is_2d, ndim
+
+
 def get_supervised_loader(
     data_paths: Tuple[str],
     raw_key: str,
@@ -108,16 +120,7 @@ def get_supervised_loader(
     Returns:
         The PyTorch dataloader.
     """
-
-    # Check for 2D or 3D training
-    try:
-        z, y, x = patch_shape
-        ndim = 2 if z == 1 else 3
-    except ValueError:
-        y, x = patch_shape
-        ndim = 2
-    print("ndim is: ", ndim)
-
+    _, ndim = determine_ndim(patch_shape)
     if label_transform is not None:  # A specific label transform was passed, do nothing.
         pass
     elif add_boundary_transform:
@@ -166,7 +169,7 @@ def supervised_training(
     val_paths: Tuple[str],
     label_key: str,
     patch_shape: Tuple[int, int, int],
-    save_root: str,
+    save_root: Optional[str] = None,
     raw_key: str = "raw",
     batch_size: int = 1,
     lr: float = 1e-4,
@@ -236,14 +239,7 @@ def supervised_training(
         check_loader(val_loader, n_samples=4)
         return
 
-    # Check for 2D or 3D training
-    try:
-        z, y, x = patch_shape
-    except ValueError:
-        y, x = patch_shape
-        z = 1
-    is_2d = z == 1
-
+    is_2d, _ = determine_ndim(patch_shape)
     if is_2d:
         model = get_2d_model(out_channels=out_channels)
     else: