unable to train a segmentation model

[qwang0225]

Hi, I am a physician new to medical imaging research. I've been working on training a semantic segmentation model to detect bleeding in CT head scans. My dataset comprises about 70 CT images and corresponding mask files. The images are 3D pixel arrays sized at 512x512xHeight, where 'Height' varies from 30 to 50 slices per scan. The masks mirror this shape but only contain binary values (0 or 1), with '1' indicating bleeding pixels.

The challenge is the significant imbalance in the pixel values in the masks, where over 99% are '0', and some samples completely lack positive ('1') pixels. This seems typical for minor bleeding in CT head imaging. To address this imbalance, I've applied data augmentation and focal loss techniques, adjusting weights. Nothing worked.

The model continued to underperform, with the mean Dice coefficient lingering around 0.1 even after training for hundreds of epochs.
I hope some of you with more experience could help out. Thanks.

Here is the code, sorry it's kind of long.

images = sorted(glob.glob(os.path.join(root, "ct_scans", "*.nii")))
labels = sorted(glob.glob(os.path.join(root, "masks", "*.nii")))

data_dicts = [{"image": image_name, "label": label_name} for image_name, label_name in zip(images, labels)]
train_files, val_files = data_dicts[8:], data_dicts[:8]

train_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        ScaleIntensityRanged(
            keys=["image"],
            a_min=-20,
            a_max=100,
            b_min=0.0,
            b_max=1.0,
            clip=True,
        ),
        ScaleIntensityRanged(
            keys=["label"],
            a_min=0.0, a_max=255.0,
            b_min=0.0, b_max=1.0,  # New intensity range
            clip=True,
        ),
        #CropForegroundd(keys=["image", "label"], source_key="image"),
        #Orientationd(keys=["image", "label"], axcodes="RAS"),
        Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 2.0), mode=("bilinear", "nearest")),
        RandCropByPosNegLabeld(
            keys=["image", "label"],
            label_key="label",
            spatial_size=(64, 64, 24),
            pos=1,
            neg=1,
            num_samples=8,
            image_key="image",
            image_threshold=0,
        ),
        #CropForegroundd(keys=["image", "label"], source_key="label", roi_center=(32, 32, 12), roi_size=(64, 64, 24)),
        #ToTensord(keys=["image", "label"]),
        # user can also add other random transforms
         RandAffined(
             keys=['image', 'label'],
             mode=('bilinear', 'nearest'),
             prob=1.0, spatial_size=(64, 64, 24),
             rotate_range=(0, 0, np.pi/15),
             scale_range=(0.1, 0.1, 0.1)),
    ]
)
val_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        ScaleIntensityRanged(
            keys=["image"],
            a_min=-20,
            a_max=100,
            b_min=0.0,
            b_max=1.0,
            clip=True,
        ),
        ScaleIntensityRanged(
            keys=["label"],
            a_min=0.0, a_max=255.0,
            b_min=0.0, b_max=1.0,  # New intensity range
            clip=True,
        ),
        #CropForegroundd(keys=["image", "label"], source_key="image"),
        #Orientationd(keys=["image", "label"], axcodes="RAS"),
        Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 2.0), mode=("bilinear", "nearest")),
        #ToTensord(keys=["image", "label"]),
    ]
)

train_ds = CacheDataset(data=train_files, transform=train_transforms, cache_rate=1.0, num_workers=4)
train_loader = DataLoader(train_ds, batch_size=4, shuffle=True, num_workers=4)
val_ds = CacheDataset(data=val_files, transform=val_transforms, cache_rate=1.0, num_workers=4)
val_loader = DataLoader(val_ds, batch_size=1, num_workers=4)

class FocalLoss(nn.Module):
    def __init__(self, alpha=0.5, gamma=2.0, reduction='mean'):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.reduction = reduction

    def forward(self, inputs, targets):
        # Convert inputs from [N, C, D, H, W] to [N, D, H, W] using sigmoid
        # Assuming binary classification (C=2), and targets are [N, 1, D, H, W]
        inputs = torch.sigmoid(inputs[:, 1, ...])  # Taking the second channel for the foreground class

        # Compute Binary Cross Entropy
        BCE_loss = F.binary_cross_entropy(inputs, targets.squeeze(1), reduction='none')
        pt = torch.exp(-BCE_loss)  # Prevents nans
        F_loss = self.alpha * (1 - pt) ** self.gamma * BCE_loss

        # Apply reduction
        if self.reduction == 'mean':
            return torch.mean(F_loss)
        elif self.reduction == 'sum':
            return torch.sum(F_loss)
        else:
            return F_loss


class GDiceLoss(nn.Module):
    def __init__(self, weights=None, alpha=0.5, gamma=2.0, reduction='mean', **kwargs):
        super(GDiceLoss, self).__init__()
        self.dice_loss = DiceLoss(to_onehot_y=True, softmax=True, **kwargs)
        self.focal_loss = FocalLoss(alpha=alpha, gamma=gamma, reduction=reduction)
        self.weights = weights

    def forward(self, input, target):
        dice_loss = self.dice_loss(input, target)
        focal_loss = self.focal_loss(input, target)

        # Apply class weights to Dice Loss if provided
        if self.weights is not None:
            dice_loss = dice_loss * self.weights

        # Combine Dice Loss and Focal Loss
        total_loss = dice_loss.mean() + focal_loss
        return total_loss

device = torch.device("cuda:0")
#device = "cpu"
model = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 32, 64, 128),
    strides=(2, 2, 2),
    num_res_units=2,
    norm=Norm.BATCH,
).to(device)
weights = torch.tensor([1.0, 20.0]).to(device)
loss_function = GDiceLoss(weights=weights, alpha=0.5, gamma=2.0, reduction='mean')
optimizer = torch.optim.Adam(model.parameters(), 1e-3)
dice_metric = DiceMetric(include_background=False, reduction="mean")

# training loop
max_epochs = 200
val_interval = 2
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = []
metric_values = []
post_pred = Compose([AsDiscrete(argmax=True, to_onehot=2)])
post_label = Compose([AsDiscrete(to_onehot=2)])

# path to save best model
dir = r"C:\Users\frank\projects\ich_segmentation"

for epoch in range(max_epochs):
    print("-" * 10)
    print(f"epoch {epoch + 1}/{max_epochs}")
    model.train()
    epoch_loss = 0
    step = 0
    for batch_data in train_loader:
        step += 1
        inputs, labels = (
            batch_data["image"].to(device),
            batch_data["label"].to(device),
        )
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
        #print(outputs.shape)
        #print(outputs[0][0])
        #print(outputs[1][1])
        print(f"{step}/{len(train_ds) // train_loader.batch_size}, " f"train_loss: {loss.item():.4f}")
    epoch_loss /= step
    epoch_loss_values.append(epoch_loss)
    print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")
    if (epoch + 1) % val_interval == 0:
        model.eval()
        with torch.no_grad():
            for val_data in val_loader:
                val_inputs, val_labels = (
                    val_data["image"].to(device),
                    val_data["label"].to(device),
                )
                roi_size = (64, 64, 24)
                sw_batch_size = 4
                val_outputs = sliding_window_inference(val_inputs, 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)]
                # compute metric for current iteration
                val_outputs = [i.to(device) for i in val_outputs]
                val_labels = [i.to(device) for i in val_labels]
                dice_metric(y_pred=val_outputs, y=val_labels)

            # aggregate the final mean dice result
            metric = dice_metric.aggregate().item()
            # reset the status for next validation round
            dice_metric.reset()

            metric_values.append(metric)
            if metric > best_metric:
                best_metric = metric
                best_metric_epoch = epoch + 1
                torch.save(model.state_dict(), os.path.join(dir, "best_metric_model.pth"))
                print("saved new best metric model")
            print(
                f"current epoch: {epoch + 1} current mean dice: {metric:.4f}"
                f"\nbest mean dice: {best_metric:.4f} "
                f"at epoch: {best_metric_epoch}"
            )

[qwang0225]

It's weird that after hundreds of epochs, the output basically look very similar to the input.

[KumoLiu]

Hi @qwang0225, I noticed that you used ScaleIntensityRanged on labels. The label should be binarized, and no more intensity transforms should be applied.
Perhaps you could refer to the segmentation tutorials in the tutorial repo.

Hope it helps, thanks!

[qwang0225]

Hi, Because the positive pixels in original labels are 255, so I had to adjust this to the range of 0-1 to match the images.

[KumoLiu]

ok, got it. Just want to make sure that your label doesn't contain other float numbers.

• BTW, I noticed that you used RandCropByPosNegLabeld since your foreground is significantly smaller than your background, so you could set the pos larger, then it will sample more foreground patches, otherwise, there may be half patches only containing zero.
• I also suggest you remove RandAffined, just let the pipeline simple. When the network can be overfitted, add more augmentations.
• When using Adam, you could set the learning rate smaller, such as 1e-4 or 1e-5.

Hope it helps, thanks!

[qwang0225]

Hi, I tried different pos in RandCropByPosNegLabeld, ranging from 5 to 100. I also tried different learning rate and loss weights, setinclude_background in DiceMetric to true instead of false since my ground truth is binary labels, which did make mean dice reach up to 0.7. However, I still cannot reach a proper output.

train_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        ScaleIntensityRanged(
            keys=["image"],
            a_min=-20,
            a_max=100,
            b_min=0.0,
            b_max=1.0,
            clip=True,
        ),
            ScaleIntensityRanged(
            keys=["label"],
            a_min=0.0, a_max=255.0,
            b_min=0.0, b_max=1.0, 
            clip=True,
        ),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 2.0), mode=("bilinear", "nearest")),
        RandCropByPosNegLabeld(
            keys=["image", "label"],
            label_key="label",
            spatial_size=(64, 64, 64),
            pos=15,
            neg=1,
            num_samples=4,
            image_key="image",
            image_threshold=0,
        ),
     ]
)
val_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
         ScaleIntensityRanged(
            keys=["image"],
            a_min=-20,
            a_max=100,
            b_min=0.0,
            b_max=1.0,
            clip=True,
        ),
         ScaleIntensityRanged(
            keys=["label"],
            a_min=0.0, a_max=255.0,
            b_min=0.0, b_max=1.0, 
            clip=True,
        ),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 2.0), mode=("bilinear", "nearest")),
    ]
)

device = torch.device("cuda:0")
#device = "cpu"
model = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
    norm=Norm.BATCH,
).to(device)

weight = torch.tensor([0.2, 0.8]).to(device)
loss_function = DiceLoss(to_onehot_y=True, softmax=True, weight=weight)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=300)
dice_metric = DiceMetric(include_background=True, reduction="mean")
post_pred = Compose([EnsureType(), AsDiscrete(argmax=True, to_onehot=2)])
post_label = Compose([EnsureType(), AsDiscrete(to_onehot=2)])

Is there anything else I could be missing? Thank you so much!