Multilabel Semantic Segmentation - Training in Clara 3.1

Hello,

I successfully did trainings to achieve binary semantic segmentation (classify pixels with label 0,1), but I am now adding additional classes.

In my binary training config, ConvertToMultiChannelBasedOnBratsClasses was used (I based it on an example BRATS config). For this new training, I attempted to use ConvertToMultiChannelBasedOnLabel, however it threw the following error:
ValueError: Cannot feed value of shape (1, 3, 288, 256, 256) for Tensor ‘NV_LABEL_INPUT:0’, which has shape ‘(?, 1, 288, 256, 256)’

To attempt to fix this, I changed my output_label_channels in the image_pipeline to 3. When I tried this, I got the following error:
ValueError: Can not squeeze dim[1], expected a dimension of 1, got 3 for ‘Squeeze’ (op: ‘Squeeze’) with input shapes: [?,3,288,256,256].
which seems to be related to the dice loss calculation.

By removing ConvertToMultiChannelBasedOnLabel, and setting output_label_channels in the image_pipeline to 1 the network begins training, but I can’t help but feel something is wrong.

Any insight would be appreciated.

If you are interested, here is my entire training configuration:
{
“epochs”: 1260,

“num_interval_per_valid”: 20,
“learning_rate”: 2e-4,
“multi_gpu”: false,
“amp”: false,
“train”: {
“loss”: {
“name”: “FocalDice”,
“args”: {
“squared_pred”: true,
“is_onehot_targets”: false,
“skip_background”: true
}
},
“optimizer”: {
“name”: “Adam”,
“args”: {
“epsilon”: 1e-4
}
},
“lr_policy”: {
“name”: “DecayOnStep”,
“args”: {
“decay_ratio”: 0.1,
“decay_freq”: 50000
}
},
“model”: {
“name”: “SegResnet”,
“args”: {
“num_classes”: 3,
“use_vae”: false,
“blocks_down”: “1,2,2,4”,
“blocks_up”: “1,1,1”,
“init_filters”: 16,
“use_group_norm”: true,
“use_group_normG”: 8,
“reg_weight”: 1e-5,
“dropout_prob”: 0.2,
“final_activation”: “softmax”
}
},
“pre_transforms”: [
{
“name”: “LoadNifti”,
“args”: {
“fields”: [
“image”,
“label”
]
}
},
{
“name”: “ConvertToChannelsFirst”,
“args”: {
“fields”: [
“image”,
“label”
]
}
},
{
“name”: “FastPosNegRatioCropROI”,
“args”: {
“size”: [
288,
256,
256
],
“image_field”: “image”,
“label_field”: “label”,
“rotation”: false,
“rotation_degree”: 90,
“scale”: true,
“scale_factor”: 0.2
}
},
{
“name”: “NormalizeNonzeroIntensities”,
“args”: {
“fields”: “image”
}
}
],
“image_pipeline”: {
“name”: “SegmentationImagePipeline”,
“args”: {
“data_list_file_path”: “{DATASET_JSON}”,
“data_file_base_dir”: “{DATA_ROOT}”,
“data_list_key”: “training”,
“output_crop_size”: [
288,
256,
256
],
“output_batch_size”: 1,
“output_image_channels”: 1,
“output_label_channels”: 1,
“num_workers”: 4,
“prefetch_size”: 0
}
},

"metrics": [
  {
    "name": "ComputeAverageDice",
    "args": {
      "name": "mean_dice",
      "is_key_metric": true,
      "field": "model",
      "label_field": "label"
    }
  }
]

},
“validate”: {
“metrics”: [
{
“name”: “ComputeAverageDice”,
“args”: {
“name”: “mean_dice”,
“is_key_metric”: true,
“field”: “model”,
“label_field”: “label”
}
}
],
“image_pipeline”: {
“name”: “SegmentationImagePipeline”,
“args”: {
“data_list_file_path”: “{DATASET_JSON}”,
“data_file_base_dir”: “{DATA_ROOT}”,
“data_list_key”: “validation”,
“output_crop_size”: [
288,
256,
256
],
“output_batch_size”: 1,
“output_image_channels”: 1,
“output_label_channels”: 1,
“num_workers”: 4,
“prefetch_size”: 0
}
},
“pre_transforms”: [
{
“name”: “LoadNifti”,
“args”: {
“fields”: [
“image”,
“label”
]
}
},
{
“name”: “ConvertToChannelsFirst”,
“args”: {
“fields”: [
“image”,
“label”
]
}
},
{
“name”: “FastPosNegRatioCropROI”,
“args”: {
“size”: [
288,
256,
256
],
“image_field”: “image”,
“label_field”: “label”,
“rotation”: false,
“rotation_degree”: 90,
“scale”: true,
“scale_factor”: 0.2
}
},
{
“name”: “NormalizeNonzeroIntensities”,
“args”: {
“fields”: “image”
}
}
],
“inferer”: {
“name”: “TFSimpleInferer”
}
}
}