Tuning new CNN module for specific dataset

TARDIS has general training and prediction model build in. With TARDIS you can train your UNet, ResNet, UNet3Plus and our custom build FNet CNN model.

Training TARDIS CNN models

1. Data preparation

TARDIS requires a folder containing training datasets as input. The dataset should include images in formats such as .tif, .mrc, .rec, or .am, and corresponding binary masks/labels in .tif, .mrc, .rec, .am, or .csv formats. The images and mask files must share the same names, with mask files having an additional _mask prefix.

E.g. For example image.mrc and image_mask.mrc or image_mask.csv

Tips:
  • For .mrc files:

    Ensure your data is not flipped. You can verify this by loading your file in Fiji and checking if the image and mask file overlap.

  • Normalization:

    No normalization is required. It is best to use raw image files, as TARDIS will handle all necessary normalizations.

2. Training

To train a CNN model with TARDIS, use the following command to get detailed information:

tardis_cnn_train --help

Basic usage:

To run training. To start training, use the following command:

tardis_cnn_train -dir <path-to-your-dataset; str> -ps <patch_size; int> -cnn <cnn_type; str> -b <batch_size> -cs 3gcl <or 2gcl for 2D>

TARDIS will create a folder named cnn_model_checkpoint, containing files named *_checkpoint.pth and model_weights.pth.

  • For re-training the model, use the *_checkpoint.pth file.

  • For predictions with TARDIS-em, use the model_weights.pth file.

Advance usage:

Below you can find all available arguments you can use with tardis_mt, with the explanation for their functionality:

-dir or --path: Directory path with all dataset and labels.

  • default: Current command line directory.

  • type: str

-ps or --patch_size: Image size used for prediction.

  • default: 64

  • type: int

-ms or --mask_size: Size of drawn mask in A. If you are using .csv files as labels.

  • default: 150

  • type: int

-cnn or --cnn_type: Type of NN used for training.

  • default: fnet_attn

  • type: str

  • options: unet, resunet, unet3plus, big_unet, fnet, fnet_attn

-co or --cnn_out_channel: Number of output channels for the NN.

  • default: 1

  • type: int

-b or --training_batch_size: Batch size.

  • default: 25

  • type: int

-cl or --cnn_layers: Number of convolution layer for NN.

  • default: 5

  • type: int

-cm or --cnn_scaler: Convolution multiplayer for CNN layers.

  • default: 32

  • type: int

-cs or --cnn_structure: Define structure of the convolution layer.

  • default: 3gcl

  • type: str

  • options: 2 or 3 - dimension in 2D or 3D; c - convolution; g - group normalization; b - batch normalization; r - ReLU; l - LeakyReLU; e - GeLu; p - PReLu

-ck or --conv_kernel: Kernel size for 2D or 3D convolution.

  • default: 3

  • type: int

-cp or --conv_padding: Padding size for convolution.

  • default: 1

  • type: int

-cmpk or --pool_kernel: Max_pooling kernel.

  • default: 2

  • type: int

-l or --cnn_loss: Loss function use for training.

  • default: BCELoss

  • type: str

  • options: AdaptiveDiceLoss, BCELoss, WBCELoss, BCEDiceLoss, CELoss, DiceLoss, ClDiceLoss, ClBCELoss, SigmoidFocalLoss, LaplacianEigenmapsLoss, BCEMSELoss

-lr or --loss_lr_rate: Learning rate for NN.

  • default: 0.0005

  • type: float

-lrs or --lr_rate_schedule: If True learning rate scheduler is used.

  • default: False

  • type: bool

-dv or --device: Define which device use for training:

  • default: 0

  • type: str

  • options: gpu - Use ID 0 gpus; cpu - Usa CPU; mps - Apple silicon; 0-9 - specified gpu device id to use

-w or --warmup: Number of warmup steps.

  • default: 100

  • type: int

-e or --epochs: Number of epoches.

  • default: 10000

  • type: int

-es or --early_stop: Number of epoches without improvement after which early stop is initiated. Default should is 10% of the total number of epochs.

  • default: 1000

  • type: int

-cch or --cnn_checkpoint: If indicated, dir to training checkpoint to reinitialized training.

  • default: None

  • type: str

-dp or --dropout_rate: If indicated, value of dropout for CNN.

  • default: 0.5

  • type: float

3.1 Pre-train model from scratch

To run re-training:

tardis_cnn_train -dir <path-to-your-dataset; str> -ps <patch_size; int> -cnn <cnn_type; str> -b <batch_size> -cch <checkpoint.pth_file_dir>

3.2 Fine-tune existing models

All TARDIS models are stored locally in ~/tardis_em/

For example a default model for membrane segmentation can be found in

./tardis_em/fnet_attn_32/membrane_3d/model_weights.pth

In order to fine-tune it on your existing data:

tardis_cnn_train ... -cch ./tardis_em/fnet_attn_32/membrane_3d/model_weights.pth

4. Predict with train model

To predict images with your newly train model, you can use the following command:

tardis_predict --help
tardis_predict -dir <dir to folder of file to predict> -ch <model_weight.pth_directory> -ps <patch_size> -out <output_format> mrc|tif|rec|am -rt True -ct <CNN_threshold> -dv 0