train_config_regression.yaml

# model configuration
model:
  # model class, e.g. UNet3D, ResidualUNet3D
  name: UNet3D
  # number of input channels to the model
  in_channels: 1
  # number of output channels
  out_channels: 1
  # determines the order of operators in a single layer (gcr - GroupNorm+Conv3d+ReLU)
  layer_order: gcr
  # number of features at each level of the U-Net
  f_maps: [16, 32, 64, 128, 256]
  # number of groups in the groupnorm
  num_groups: 8
  # if True applies the final normalization layer (sigmoid or softmax), otherwise the networks returns the output from the final convolution layer; use False for regression problems, e.g. de-noising
  is_segmentation: false
# trainer configuration
trainer:
  # path to the checkpoint directory
  checkpoint_dir: "CHECKPOINT_DIR"
  # path to latest checkpoint; if provided the training will be resumed from that checkpoint
  resume: null
  # path to the best_checkpoint.pytorch; to be used for fine-tuning the model with additional ground truth
  # make sure to decrease the learning rate in the optimizer config accordingly
  pre_trained: null
  # how many iterations between validations
  validate_after_iters: 200
  # how many iterations between tensorboard logging
  log_after_iters: 100
  # max number of epoc0hs
  max_num_epochs: 100
  # max number of iterations
  max_num_iterations: 100000
  # model with higher eval score is considered better
  eval_score_higher_is_better: true
# optimizer configuration
optimizer:
  # initial learning rate
  learning_rate: 0.0002
  # weight decay
  weight_decay: 0.00001
# loss function configuration
loss:
  # loss function to be used during training
  name: SmoothL1Loss
  # a target value that is ignored and does not contribute to the input gradient
  ignore_index: null
# evaluation metric configuration
eval_metric:
  # peak signal to noise ration
  name: PSNR
  # a target label that is ignored during metric evaluation
  ignore_index: null
# learning rate scheduler configuration
lr_scheduler:
  # reduce learning rate when evaluation metric plateaus
  name: ReduceLROnPlateau
  # use 'max' if eval_score_higher_is_better=True, 'min' otherwise
  mode: max
  # factor by which learning rate will be reduced
  factor: 0.1
  # number of *validation runs* with no improvement after which learning rate will be reduced
  patience: 10
# data loaders configuration
loaders:
  # class of the HDF5 dataset, currently StandardHDF5Dataset and LazyHDF5Dataset are supported.
  dataset: StandardHDF5Dataset
  # batch dimension; if number of GPUs is N > 1, then a batch_size of N * batch_size will automatically be taken for DataParallel
  batch_size: 1
  # how many subprocesses to use for data loading
  num_workers: 4
  # path to the raw data within the H5
  raw_internal_path: raw
  # path to the the label data within the H5
  label_internal_path: random
  # path to the pixel-wise weight map withing the H5 if present
  weight_internal_path: null
  # configuration of the train loader
  train:
    # absolute paths to the training datasets; if a given path is a directory all H5 files ('*.h5', '*.hdf', '*.hdf5', '*.hd5')
    # inside this this directory will be included as well (non-recursively)
    file_paths:
        - "PATH_TO_TRAIN_SET"

    # SliceBuilder configuration, i.e. how to iterate over the input volume patch-by-patch
    slice_builder:
      # SliceBuilder class
      name: SliceBuilder
      # train patch size given to the network (adapt to fit in your GPU mem, generally the bigger patch the better)
      patch_shape: [128, 128, 128]
      # train stride between patches
      stride_shape: [32, 32, 32]

    # data transformations/augmentations
    transformer:
      raw:
        # apply min-max scaling and map the input to [-1, 1]
        - name: Normalize
        - name: RandomFlip
        - name: RandomRotate90
        - name: RandomRotate
          # rotate only in ZY only since most volumetric data is anisotropic
          axes: [[2, 1]]
          angle_spectrum: 30
          mode: reflect
        - name: ToTensor
          expand_dims: true
      label:
        # apply min-max scaling and map the input to [-1, 1]
        - name: Normalize
        - name: RandomFlip
        - name: RandomRotate90
        - name: RandomRotate
          # rotate only in ZY only since most volumetric data is anisotropic
          axes: [[2, 1]]
          angle_spectrum: 30
          mode: reflect
        - name: ToTensor
          expand_dims: true

  # configuration of the validation loaders
  val:
    # paths to the validation datasets; if a given path is a directory all H5 files ('*.h5', '*.hdf', '*.hdf5', '*.hd5')
    # inside this this directory will be included as well (non-recursively)
    file_paths:
        - "PATH_TO_VAL_SET"

    # SliceBuilder configuration
    slice_builder:
      # SliceBuilder class
      name: SliceBuilder
      # validation patch (can be bigger than train patch since there is no backprop)
      patch_shape: [128, 128, 128]
      # validation stride (validation patches doesn't need to overlap)
      stride_shape: [128, 128, 128]

    # no data augmentation during validation
    transformer:
      raw:
        # apply min-max scaling and map the input to [-1, 1]
        - name: Normalize
        - name: ToTensor
          expand_dims: true
      label:
        # apply min-max scaling and map the input to [-1, 1]
        - name: Normalize
        - name: ToTensor
          expand_dims: true