Training of sarcomere U-Net++

This tutorial explains the training of a U-Net++ neural network model for the simultaneous prediction of sarcomere Z-bands, M-bands, sarcomere orientation maps, sarcomere masks and cell masks from microscopy images. SarcAsM uses our bio-image-unet package, see https://github.com/danihae/bio-image-unet. We strongly recommend using a GPU-equipped workstation or server for training. Make sure that CUDA toolkit along the respective version of PyTorch are installed and verify the correct installation by

import torch
torch.cuda.is_available()

The generation of a training data set is explained in this tutorial.

For training, the training data needs to be in the following directory structure, with exactly identical filenames for images and targets:

training_data/
|
├── image
│   ├── 105.tif
│   ├── 111.tif
│   ├── image123xyz.tif
│   ├── 121.tif
│   ├── 1.tif
│   ├── 2.tif
│   └── 83.tif
├── zbands
│   ├── 105.tif
│   ├── 111.tif
│   ├── image123xyz.tif
│   ├── 121.tif
│   ├── 1.tif
│   ├── 2.tif
│   └── 83.tif
├── orientation
│   ├── 105.tif
│   ├── 111.tif
│   ├── image123xyz.tif
│   ├── 121.tif
│   ├── 1.tif
│   ├── 2.tif
│   └── 83.tif
└── ....
    ├── ....

# set paths of directories for images and targets
dir_training_data = '../../training_data_tutorial/'  # modify path if needed
dir_images = dir_training_data + '/image/'
dir_zbands = dir_training_data + '/zbands/'
dir_zbands_prelim = dir_training_data + '/zbands_prelim/'
dir_mbands = dir_training_data + '/mbands/'
dir_orientation = dir_training_data + '/orientation/'
dir_sarcomere_mask = dir_training_data + '/sarcomere_mask/'
dir_cell_mask = dir_training_data + '/cell_mask/'

Prepare and process training data

Prior to training, the training images and labels are processed and augmented. For the different options for processing and augmentation (add noise, blur, adjust contrast, …) see docstring or API reference.

import os
import bio_image_unet.multi_output_unet as unet

# temp dir
dir_temp = '../../training_temp/'  # adjust when necessary
os.makedirs(dir_temp, exist_ok=True)

# define targets
target_dirs = [dir_zbands, dir_mbands, dir_orientation, dir_sarcomere_mask, dir_cell_mask]
target_types = {'zbands': 'mask', 'mbands': 'mask', 'orientation': 'mask', 'sarcomere_mask': 'mask', 'cell_mask': 'mask'}

# prepare and augment training data
data = unet.DataProcess(image_dir=dir_images, target_dirs=target_dirs, target_types=target_types, aug_factor=6, data_dir=dir_temp)

Set training parameters and train

For different training parameters, check the docstring print(unet.Trainer.__doc__) or API reference.

from bio_image_unet.multi_output_unet.multi_output_nested_unet import MultiOutputNestedUNet_3Levels

# prepare output heads
output_heads = {'zbands': {'channels': 1, 'activation': 'sigmoid', 'loss': 'BCEDiceLoss', 'weight': 2},
                'midlines': {'channels': 1, 'activation': 'sigmoid', 'loss': 'BCEDiceLoss', 'weight': 2},
                'cell_mask': {'channels': 1, 'activation': 'sigmoid', 'loss': 'BCEDiceLoss', 'weight': 0.5},
                'sarcomere_mask': {'channels': 1, 'activation': 'sigmoid', 'loss': 'BCEDiceLoss', 'weight': 1},
                'orientation': {'channels': 2, 'activation': None, 'loss': 'VectorFieldLoss', 'weight': 1}
                }

trainer = unet.Trainer(data, network=MultiOutputNestedUNet_3Levels, batch_size=4, num_epochs=100, lr=1e-4, n_filter=64,
                  save_iter=True, output_heads=output_heads, deep_supervision=True, load_weights=False, levels=3,
                  save_dir='../../training_temp/')

# start training
trainer.start()

After training is completed, the model parameters model.pt are stored in the save_dir.