train_unet_sem.yaml

  1mode: train
  2# contents to define a model
  3model:
  4## model architecture, SeM indicates Segmentation Model
  5  name: SeM
  6  ## loss function of architecture. For SeM, you need to specify the segmentation loss.
  7  segmentation_losses:
  8    - name: Dice
  9      weight: 1.0
 10    - name: BCEL
 11      weight: 1.0
 12  ## define encoder and decoder
 13  encoder:
 14    ### encoder name
 15    name: UNet
 16    ### input image size (after augmented)
 17    image_size: [320, 256]
 18    ### number of input image channels
 19    in_channel: 3
 20    ### number of output channels, for segmentation, it should be the number of categories
 21    out_channel: 1
 22    ### number of channel for image patches
 23    patch_channel: 32
 24    ### size of image patch, for 2D image, 'patch size = 2' indicates a 2*2 image patch.
 25    patch_size: 2
 26    ### number of channel for each layer in down-sample layers.
 27    ### The length of list is the number of down-sample layers
 28    down_channels: [64, 128, 256]
 29    ### number of channel for each layer in middle layers.
 30    ### The length of list is the number of middle layers
 31    middle_channels: [512, 512]
 32    ### building block
 33    building_block: conv
 34    ### normalization
 35    normalization: batch
 36# data loader for training
 37loader:
 38  ## define a dataset
 39  dataset:
 40    name: ImgSegDataset
 41    data_dir: images
 42    label_dir: masks
 43    data_suffix: jpg
 44  ## merge listed datasets into one
 45  data_path_list:
 46    - path/to/datasets/CVC-ClinicDB/fold1/
 47    - path/to/datasets/CVC-ClinicDB/fold2/
 48    - path/to/datasets/CVC-ClinicDB/fold3/
 49  ### other parameters related to dataloader
 50  ### refer to https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader
 51  batch_size: 16
 52  num_workers: 8
 53  shuffle: true
 54  ## augmentations performed on each sample from the dataset
 55  data_transforms:
 56    - name: Resize
 57      size: [320, 256]
 58    - name: ToTensor
 59# data loader for validation
 60val_loader:
 61  dataset:
 62    name: ImgSegDataset
 63    data_dir: images
 64    label_dir: masks
 65    data_suffix: jpg
 66    data_path: path/to/datasets/CVC-ClinicDB/fold4/
 67  batch_size: 8
 68  num_workers: 8
 69  shuffle: false
 70  data_transforms:
 71    - name: Resize
 72      size: [320, 256]
 73    - name: ToTensor
 74# define a optimizer
 75optimizer:
 76  name: Adam
 77  lr: 0.0003
 78  weight_decay: 0.00000001
 79# define a learning rate scheduler
 80lr_scheduler:
 81  name: LinearLR
 82  start_factor: 1.0
 83  end_factor: 0.01
 84# evaluation metrics
 85evaluation_metrics:
 86  seg:
 87    - name: Dice
 88    - name: ACC
 89    - name: mIoU
 90
 91score_metric:
 92  name: Dice
 93  higher_is_better: true
 94
 95# max training epochs
 96max_epoch: 500
 97# in warm-up epoch, learning rate will be fixed as the initial value
 98warmup_epoch: 2
 99# write intermediate results to tensorboard for visualization
100write_after_iters: 5
101# save checkpoint
102save_after_epochs: 2
103# directory for checkpoints
104check_point_dir: path/to/checkpoint/CVC-ClinicDB/UNet