Configuration

Empanada uses YAML to define training and inference parameters.

Training

There are four sections to the training config: DATASET, MODEL, TRAIN and EVAL.

The DATASET section gives a name to the training dataset, names the classes that will be segmented and lists the ids of classes that require instance segmentation.

An example for instance segmentation of a single class:

DATASET:
  # name of dataset tracked by MLFLOW (see Logging in docs)
  dataset_name: "Mitochondria"
  # name of the segmentation class
  class_names: [ "mito" ]
  # only 1 segmentation class, id == 1
  labels: [ 1 ]
  # the 1 segmentation class has instances
  thing_list: [ 1 ]
  # pixel mean and std in dataset images
  norms: { mean: 0.508979, std: 0.148561 }

An example for true panoptic segmentation:

DATASET:
  dataset_name: "Organelles"
  # name of the segmentation class
  class_names: [ "er", "nuclei", "vesicles", "mito" ]
  # the four segmentation classes
  labels: [ 1, 2, 3, 4 ]
  # ER is semantic, all the others are instance
  thing_list: [ 2, 3, 4 ]
  norms: { mean: 0.577841, std: 0.114811 }

The MODEL section provides arguments to the model class. Any model arg or kwarg can be added to the config file and passed to the model during training.

MODEL:
  # name of the model architecture
  arch: "PanopticDeepLabPR"
  # the encoder network to use
  encoder: "resnet50"
  # kwargs for this particular model architecture
  num_classes: 1
  stage4_stride: 16
  decoder_channels: 256
  low_level_stages: [ 1 ]
  low_level_channels_project: [ 32 ]
  atrous_rates: [ 2, 4, 6 ]
  aspp_channels: null
  aspp_dropout: 0.5
  ins_decoder: True
  ins_ratio: 0.5

The TRAIN section defines all parameters used to train the model, it’s pretty extensive to allow for maximum flexibility.

TRAIN:
  # name of the run stored in MLFLOW under "dataset_name"
  run_name: "Panoptic DeepLab PR Baseline"

  # Directories containing images and masks for the segmentation
  # task. Multiple directories can be appended to the main train_dir
  train_dir: "mitochondria_data"
  additional_train_dirs: [ "new_mitochondria_data" ]

  # directory to save models to
  # and epochs between saving
  model_dir: "pdlpr_baselines/"
  save_freq: 1

  # path to .pth file for resuming training
  resume: null

  # encoder pretrained weights
  encoder_pretraining: "cem1.5m_swav_resnet50_200ep_balanced.pth.tar"
  # pretrained weights for the entire model, useful for finetuning
  whole_pretraining: null
  # layers in the encoder to finetune, choice of
  # "all", "none", "stage1", "stage2", "stage3", "stage4"
  finetune_layer: "all"

  # the learning rate schedule,
  # see torch.optim.lr_scheduler
  lr_schedule: "OneCycleLR"
  schedule_params:
    max_lr: 0.003
    epochs: 30
    steps_per_epoch: 339
    pct_start: 0.3

  # automatic mixed precision
  amp: True

  # setup the optimizer, see torch.optim
  optimizer: "AdamW"
  optimizer_params:
    weight_decay: 0.1

  # parameters to pass to the PDL loss function, see empanada.losses
  criterion_params:
    ce_weight: 1
    mse_weight: 200
    l1_weight: 0.01
    top_k_percent: 0.2
    confidence_loss: False
    cl_weight: 0.1
    pr_weight: 1

  # training performance metric to track, see empanada.metrics
  print_freq: 50
  metrics: [ "IoU" ]
  metric_params:
    topk: 1

  # dataset and dataloader parameters, see empanada.data
  batch_size: 64
  dataset_class: "SingleClassInstanceDataset"
  weight_gamma: 0.3
  workers: 8

  # augmentations to apply to images,
  # nearly any in albumentations are supported
  augmentations:
    # aug name and kwargs
    - { aug: "RandomScale", scale_limit: [ -0.9, 1 ]}
    - { aug: "PadIfNeeded", min_height: 256, min_width: 256, border_mode: 0 }
    - { aug: "RandomCrop", height: 256, width: 256}
    - { aug: "Rotate", limit: 180, border_mode: 0 }
    - { aug: "RandomBrightnessContrast", brightness_limit: 0.3, contrast_limit: 0.3 }
    - { aug: "HorizontalFlip" }
    - { aug: "VerticalFlip" }

  # parameters for multi-GPU training
  multiprocessing_distributed: False
  gpu: null
  world_size: 1
  rank: 0
  dist_url: "tcp://localhost:10001"
  dist_backend: "nccl"

The EVAL section handles the evaluation of the model during training. It’s an optional section. Setting aside a validation dataset is useful for tuning training parameters though.

EVAL:
  eval_dir: "eval_mitochondria"
  # track segmentation of particular images in the
  # evaluation data, stored in MLFLOW
  eval_track_indices: [ 10, 32, 19, 7 ]

  # how often to record segmentations
  eval_track_freq: 10

  # epochs between evaluation
  epochs_per_eval: 1

  # parameters needed for eval_metrics
  # see empanada.metrics
  metrics: [ "IoU", "PQ", "F1" ]
  metric_params:
      topk: 1
      labels: [ 1 ]
      label_divisor: 1000
      iou_thr: 0.5

  # parameters used for panoptic inference
  # see empanada.inference.engines
  engine_params:
    thing_list: [ 1 ]
    label_divisor: 1000
    stuff_area: 64
    void_label: 0
    nms_threshold: 0.1
    nms_kernel: 7
    confidence_thr: 0.5

Inference

The config file for model inference has only a single section and is less complicated than the training config.

# axes to predict for 3d
# for stack inference, just 'xy'
axes: [ 'xy', 'xz', 'yz' ]

# list of all segmentation labels
labels: [ 1 ]

# parameters for the inference engine
# see empanada.inference.engines
engine_params:
  median_kernel_size: 7
  thing_list: [ 1 ]
  label_divisor: 20000
  stuff_area: 64
  void_label: 0
  nms_threshold: 0.1
  nms_kernel: 7
  confidence_thr: 0.3
  input_scale: 1
  scales: [ 1 ]

# parameters for instance matching
# across 2d images
matcher_params:
  merge_iou_thr: 0.25
  merge_ioa_thr: 0.25
  force_connected: True

# parameters for the consensus algorithm
consensus_params:
  pixel_vote_thr: 0.5
  cluster_iou_thr: 0.75

# object filters, see empanada.inference.filters
filters:
  - { name: "remove_small_objects", min_size: 500 }
  - { name: "remove_pancakes", min_span: 4 }

Inheritance

Fortunately, defining a training config file from scratch is almost never necessary. To modify one or a few parameters just add the BASE key to the top of the config file and point it to the parent config file. Any parameters that you define in this new config file will overwrite the parameters in the BASE config, everything else gets inherited.

For example, to change the number of training epochs used by the OneCycle policy, the complete config file would look like this:

BASE: "path_to_parent_training.yaml"
TRAIN:
  # not necessary to change the run name
  # but it makes tracking experiments easier
  run_name: "Panoptic DeepLab PR Baseline - 100 Epochs"
  schedule_params:
    epochs: 100

Or the change the median filter size for inference:

BASE: "path_to_parent_inference.yaml"

engine_params:
  median_kernel_size: 3

Inheritance is a powerful way to test many different configurations with minimal effort and should be used as often as possible. An arbitrary number of levels are possible too; for example, a config file can inherit from another file that also uses inheritance.