How to prepare labels for multi-class segmentation

[hermancollin]

Hello. I am struggling to train a 2D UNet segmentation model using monai. We work with microscopy images and want the model to predict 2 classes (3 if you count the background). So far, here is the material I based my training script on:

• the 2d segmentation tutorial (tutorials/2d_segmentation/torch/unet_training_dict.py) for general outline. I use the dictionary-based code.
• the brats 3d segmentation tutorial (tutorials/3d_segmentation/brats_segmentation_3d.ipynb) which gives a basic idea of how to handle multiclass labels. Although my implementation is slightly different, from my understanding it does the same thing.
• the dice metric/loss notes (tutorials/modules/dice_loss_metric_notes.ipynb) which give a better understanding of the interface of the loss and the metrics functions.

However the training doesn't work. I'm concerned that I'm mishandling the labels?

I use a dictionary to align the inputs and the labels. Note that I use 2 different keys for the 2 labels:

train_files = [{"im": str(img), "seg-ax": str(ax), "seg-my": str(my)} for (img, ax, my) in train]
val_files = [{"im": str(img), "seg-ax": str(ax), "seg-my": str(my)} for (img, ax, my) in val]
test_files = [{"im": str(img), "seg-ax": str(ax), "seg-my": str(my)} for (img, ax, my) in test]

Training/validation transforms:

train_transforms = Compose(
    [
        LoadImaged(keys=["im", "seg-ax", "seg-my"]),
        EnsureChannelFirstd(keys=["im", "seg-ax", "seg-my"]),
        NormalizeIntensityd(keys="im"),
        RandAffined(
            keys=["im", "seg-ax", "seg-my"], 
            prob=1.0, 
            rotate_range=np.pi/64, 
            scale_range=0.05,
            translate_range=(10, 10),
            padding_mode="zeros",
            device=device
        ),
        Rand2DElasticd(
            keys=["im", "seg-ax", "seg-my"],
            prob=0.5,
            spacing=(30, 30),
            magnitude_range=(1, 2),
            padding_mode="zeros",
            device=device,
        ),
    ]
)
val_transforms = Compose(
    [
        LoadImaged(keys=["im", "seg-ax", "seg-my"]),
        EnsureChannelFirstd(keys=["im", "seg-ax", "seg-my"]),
        NormalizeIntensityd(keys="im")
    ]
)

Load the images in a preliminary Dataset, then transform it to GridPatchDataset:

train_data = monai.data.Dataset(data=train_files, transform=train_transforms)
val_data = monai.data.Dataset(data=val_files, transform=val_transforms)

patch_iterator = monai.data.PatchIterd(keys=["im", "seg-ax", "seg-my"], patch_size=(256, 256), mode='constant')
bs = 4
# need num_worker=0 in dataloader for GPU
nw = 0

# training data loader
train_ds = monai.data.GridPatchDataset(data=train_data, patch_iter=patch_iterator)
train_loader = DataLoader(train_ds, batch_size=bs, num_workers=nw, collate_fn=list_data_collate)

# validation data loader
val_ds = monai.data.GridPatchDataset(data=val_data, patch_iter=patch_iterator)
val_loader = DataLoader(val_ds, batch_size=bs, num_workers=nw, collate_fn=list_data_collate)

Define the 2-class segmentation model

device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = monai.networks.nets.UNet(
    spatial_dims=2,
    in_channels=1,
    out_channels=2,
    channels=(64, 128, 256, 512),
    strides=(1, 1, 1),
    act="ReLU",
    dropout=0.2,
    # nn.BatchNorm2d momentum defaults to 0.1
    norm="BATCH",
    # order of conv operations in ivadomed is ReLU -> batch-norm -> dropout
    adn_ordering="AND"
).to(device)

Metric and loss:

dice_metric = DiceMetric(include_background=False, reduction="mean", get_not_nans=False)
loss_function = monai.losses.DiceLoss(include_background=False, sigmoid=True)

Post-transforms:

post_pred = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
post_label = Compose([AsDiscrete()])

Inside the training loop, I stack the 2 labels together to create a single 2-channel ground-truth. I do this to match the output of the model, which has size [B, 2, H, W] for 2 output channels.

...
        for batch_data in train_loader:
            iteration += 1
            inputs = batch_data[0]["im"].to(device)
            mask_ax, mask_my = batch_data[0]["seg-ax"].to(device), batch_data[0]["seg-my"].to(device)
            # merge axon and myelin masks into a single label tensor of size (bs, 2, 256, 256)
            labels = torch.cat([mask_ax, mask_my], dim=1)

            optimizer.zero_grad()
            outputs = model(inputs)

            loss = loss_function(outputs, labels)
            loss.backward()

            optimizer.step()
...

For validation, I pretty much do the same thing:

...
            for val_data in val_loader:
                val_images = val_data[0]["im"].to(device)
                val_labels = (val_data[0]["seg-ax"].to(device), val_data[0]["seg-my"].to(device))
                val_labels = torch.cat(val_labels, dim=1)
                roi_size = (256, 256)
                sw_batch_size = 4
                val_outputs = sliding_window_inference(val_images, roi_size, sw_batch_size, model)
                val_outputs = [post_pred(i) for i in decollate_batch(val_outputs)]
                val_labels = [post_label(i) for i in decollate_batch(val_labels)]
                dice_metric(y_pred=val_outputs, y=val_labels)
            metric = dice_metric.aggregate().item()
            dice_metric.reset()
...

Would anyone have any insight on what's the problem here? The model is not able to learn at all so I must be doing something wrong with the loss function and the validation metric but I'm a bit stuck right now. Should I try to one-hot encode the predictions and labels?

[KumoLiu]

Hi @hermancollin, thanks for your interest here. Could you show me the whole error message?

[hermancollin]

Hi @KumoLiu. There is actually no error message, the training can be successfully completed but the training loss and the validation metrics are suspicious. For example, for more than 100 epochs on my data, I get these:

This is the kind of predictions I get:

pred	GT

I don't understand why it's failing this bad. Is it possible that when loss.backward() is called, the gradient is not propagated separately for both classes? Note that the validation Dice metrics approaches 2. Looking at the predictions, this is weird because I don't think the Dice metric would give a value near 1 for any of the 2 classes predicted.

[KumoLiu]

Hi @hermancollin, I have several suggestions here.

• Firstly, remove RandAffined these complexity random transforms first to get a baseline.
• Second, you should set different mode for image and segmentation. Such as mode in RandAffined. Use GridSampleMode.NEAREST instead of the default mode for segmentation to avoid non-integer pixel. I guess this one can be the biggest reason for the your poor training. But I still suggest you remove these transforms first for a baseline.
  

  MONAI/monai/transforms/spatial/array.py

  Line 2149 in 554b172

  mode: str | int = GridSampleMode.BILINEAR,

• After pre-transforms, visualize your data to check if they are make sense.

Hope these can help you, thanks!

[hermancollin]

Hi @KumoLiu. Unfortunately, I visualized my pre-transforms and everything looks normal.
For the interpolation mode, I 100% agree with you that nearest neighbord should be used for the labels and, for example, bilinear for the image itself. However, I don't think it would degrade the training up to the extent I'm showing here. I still get poor results when I ablate the Affine and Elastic transforms.

I have an intuition that the output of the model should be one-hot encoded before being fed to the loss function because that's what we usually do when dealing with multiple classes. In my case, I would need to convert the 2 output channels into a 3-value one-hot vector (1 for background, 2 for my predicted classes) but I'm not sure how to do this.

[hermancollin]

Do you think it would be possible to use a model with 3 output channels and convert the 2 masks into 3 channels (one-hot encoding) with a shared background class?