Add README for unet training

Author: constantinpape

scripts/training/README.md

@@ -0,0 +1,14 @@
+# 3D U-Net Training for Cochlea Data
+
+This folder contains the scripts for training a 3D U-Net for cell segmentation in the cochlea data.
+It contains two relevant scripts:
+- `check_training_data.py`, which visualizes the training data and annotations in napari.
+- `train_distance_unet.py`, which trains the 3D U-Net.
+
+Both scripts accept the argument `-i /path/to/data`, to specify the root folder with the training data. For example, run `python train_distance_unet.py -i /path/to/data` for training. The scripts will consider all tif files in the sub-folders of the root folder for training.
+They will load the **image data** according to the following rules:
+- Files with the ending `_annotations.tif` or `_cp_masks.tif` will not be considered as image data.
+- The other files will be considered as image data, if a corresponding file with ending `_annotations.tif` can be found. If it cannot be found the file will be excluded; the scripts will print the name of all files being excluded.
+
+The training script will save the trained model in `checkpoints/cochlea_distance_unet_<CURRENT_DATE>`, e.g. `checkpoints/cochlea_distance_unet_20250115`.
+For further options for the scripts run `python check_training_data.py -h` / `python train_distance_unet.py -h`.

scripts/training/train_distance_unet.py

@@ -6,7 +6,7 @@
 import torch_em
 from torch_em.model import UNet3d
 
-ROOT_CLUSTER = "/scratch-grete/usr/nimcpape/data/moser/lightsheet"
+ROOT_CLUSTER = "/scratch-grete/usr/nimcpape/data/moser/lightsheet/training"
 
 
 def get_image_and_label_paths(root):
@@ -77,7 +77,7 @@ def get_loader(root, split, patch_shape, batch_size, filter_empty):
     elif split == "val":
         n_samples = 20 * batch_size
 
-    sampler = torch_em.data.sampler.MinInstanceSampler(p_reject=0.95)
+    sampler = torch_em.data.sampler.MinInstanceSampler(p_reject=0.8)
     loader = torch_em.default_segmentation_loader(
         raw_paths=image_paths, raw_key=None, label_paths=label_paths, label_key=None,
         batch_size=batch_size, patch_shape=patch_shape, label_transform=label_transform,
@@ -93,6 +93,11 @@ def main():
         "--root", "-i", help="The root folder with the annotated training crops.",
         default=ROOT_CLUSTER,
     )
+    parser.add_argument(
+        "--batch_size", "-b", help="The batch size for training. Set to 8 by default."
+        "You may need to choose a smaller batch size to train on yoru GPU.",
+        default=8, type=int,
+    )
     parser.add_argument(
         "--check_loaders", "-l", action="store_true",
         help="Visualize the data loader output instead of starting a training run."
@@ -106,13 +111,13 @@ def main():
     )
     args = parser.parse_args()
     root = args.root
+    batch_size = args.batch_size
     check_loaders = args.check_loaders
     filter_empty = args.filter_empty
     run_name = datetime.now().strftime("%Y%m%d") if args.name is None else args.name
 
     # Parameters for training on A100.
     n_iterations = 1e5
-    batch_size = 8
     patch_shape = (64, 128, 128)
 
     # The U-Net.