model (3D UNet) is not learning and dice score is always below 0.10

[ankitabuntolia]

Hello,

I am new to UNet and Monai, but Monai has been so far quite useful to work on Medical Images, so thank you for maintaining and answering questions.

Problem
My mode is not learning with 100 epochs. I also tried with 200 and 400 but losses are always around 0.90 and dice score is always below 0.10

Background of task
I am working on 3D OCT Image segmentation to identify IRF, SRF, and PED fluids using Monai. My input image size is 512x496x49 and my label is also 512x496x49. An image can have 3 labels (IRF, SRF, and PED) so it's a multilabel segmentation task.
After using transform, the shape of image changes to torch.Size([1, 128, 128, 128]), label shape: torch.Size([3, 128, 128, 128]).

I am following BRAT Tutorial but the network I wanted to use UNet so I use it from Spleen Tutorial with modification discussed in #415

Here is my code

Transform

a_min=1000
a_max=30000

train_transform = Compose(
    [
      LoadImaged(keys=["image", "label"]),
      EnsureChannelFirstd(keys="image"),
      CopyItemsd(keys="image", times=3),
      SpatialCropd(keys=['image', 'label'], roi_start=(0,0,0), roi_end=(256,248,49)),
      SpatialCropd(keys=['image_1', 'label_1'], roi_start=(0,248,0), roi_end=(256,496,49),allow_missing_keys=True),
      SpatialCropd(keys=['image_2', 'label_2'], roi_start=(256,0,0), roi_end=(512,248,49),allow_missing_keys=True),
      SpatialCropd(keys=['image_3', 'label_3'], roi_start=(256,248,0), roi_end=(512,496,49),allow_missing_keys=True),
      ConcatItemsd(['image_1', 'image_2', 'image_3'], name='image', dim=1, allow_missing_keys=True),
      DeleteItemsd(['image_1', 'image_2', 'image_3']),
      ConvertToMultiChannelBasedOnOCTClassesd(keys="label"),
      ScaleIntensityRanged(keys='image', a_min=a_min, a_max=a_max, b_min=0.0, b_max=1.0, clip=False),
      Orientationd(keys=["image", "label"], axcodes="RAS"),
      # Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 5.0), mode=("bilinear", "nearest")),
      RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
      RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
      RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
      Rand3DElasticd(keys=["image", "label"], mode=("bilinear", "nearest"), prob=1.0, sigma_range=(5, 8), magnitude_range=(100, 200),
                    spatial_size=(128, 128, 128), translate_range=(-140, 140), rotate_range=(np.pi / 36, np.pi / 36, np.pi), scale_range=(0.15, 0.15, 0.15), padding_mode="border"),
      # CropForegroundd(keys=["image", "label"], source_key="label", select_fn = 6000),
      # RandCropByPosNegLabeld(keys=['image','label'], label_key='label',spatial_size = (128,124,25), pos = 1.0, neg=1.0, num_samples = 4, image_key='image', image_threshold=8000),
      NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
      EnsureTyped(keys=["image", "label"]),
    ]
)
val_transform = Compose(
    [
      LoadImaged(keys=["image", "label"]),
      EnsureChannelFirstd(keys="image"),
      ConvertToMultiChannelBasedOnOCTClassesd(keys="label"),
      ScaleIntensityRanged(keys='image', a_min=a_min, a_max=a_max, b_min=0.0, b_max=1.0, clip=False),
      Orientationd(keys=["image", "label"], axcodes="RAS"),
      # Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 5.0), mode=("bilinear", "nearest")),
      # CropForegroundd(keys=["image", "label"], source_key="label", select_fn = 6000),
      NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
      EnsureTyped(keys=["image", "label"]),
    ]
)

Dataloader

train_ds = CacheDataset(data=train_files, transform=train_transform,cache_rate=1.0, num_workers=4)
train_loader = DataLoader(train_ds, batch_size=2, shuffle=True, num_workers=0,collate_fn=pad_list_data_collate)

val_ds = CacheDataset(data=val_files, transform=val_transform, cache_rate=1.0, num_workers=4)
val_loader = DataLoader(val_ds, batch_size=1, num_workers=4)

Model

max_epochs = 100
val_interval = 1
VAL_AMP = True

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

loss_function = DiceLoss(smooth_nr=0, smooth_dr=1e-5, squared_pred=True, to_onehot_y=False, sigmoid=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-2, weight_decay=1e-5)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_epochs)

dice_metric = DiceMetric(include_background=True, reduction="mean")
dice_metric_batch = DiceMetric(include_background=True, reduction="mean_batch")

post_trans = Compose(
    [EnsureType(), Activations(sigmoid=True), AsDiscrete(threshold=0.5)]
)

# define inference method
def inference(input):
    def _compute(input):
        return sliding_window_inference(
            inputs=input,
            roi_size=(128, 128, 128),
            sw_batch_size=4,
            predictor=model,
            overlap=0.5,
        )
    if VAL_AMP:
        with torch.cuda.amp.autocast():
            return _compute(input)
    else:
        return _compute(input)


# use amp to accelerate training
scaler = torch.cuda.amp.GradScaler()
# enable cuDNN benchmark
torch.backends.cudnn.benchmark = True

root_dir = "/content/drive/MyDrive/Practical_work/out/20220327_V1_epoch100"
if not os.path.exists(root_dir):
  os.makedirs(root_dir)

print(root_dir)

Actual Training

best_metric = -1
best_metric_epoch = -1
best_metrics_epochs_and_time = [[], [], []]
epoch_loss_values = []
metric_values = []
metric_values_f1 = []
metric_values_f2 = []
metric_values_f3 = []

total_start = time.time()
for epoch in range(max_epochs):
    epoch_start = time.time()
    print("-" * 10)
    print(f"epoch {epoch + 1}/{max_epochs}")
    model.train()
    epoch_loss = 0
    step = 0
    for batch_data in train_loader:
        step_start = time.time()
        step += 1
        inputs, labels = (
            batch_data["image"].to(device),
            batch_data["label"].to(device),
        )
        optimizer.zero_grad()
        with torch.cuda.amp.autocast():
            outputs = model(inputs)
            loss = loss_function(outputs, labels)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        epoch_loss += loss.item()
        print(
            f"{step}/{len(train_ds) // train_loader.batch_size}"
            f", train_loss: {loss.item():.4f}"
            f", step time: {(time.time() - step_start):.4f}"
        )
    lr_scheduler.step()
    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),
                )
                val_outputs = inference(val_inputs)
                val_outputs = [post_trans(i) for i in decollate_batch(val_outputs)]
                dice_metric(y_pred=val_outputs, y=val_labels)
                dice_metric_batch(y_pred=val_outputs, y=val_labels)

            metric = dice_metric.aggregate().item()
            metric_values.append(metric)
            metric_batch = dice_metric_batch.aggregate()
            metric_f1 = metric_batch[0].item()
            metric_values_f1.append(metric_f1)
            metric_f2 = metric_batch[1].item()
            metric_values_f2.append(metric_f2)
            metric_f3 = metric_batch[2].item()
            metric_values_f3.append(metric_f3)
            dice_metric.reset()
            dice_metric_batch.reset()

            if metric > best_metric:
                best_metric = metric
                best_metric_epoch = epoch + 1
                best_metrics_epochs_and_time[0].append(best_metric)
                best_metrics_epochs_and_time[1].append(best_metric_epoch)
                best_metrics_epochs_and_time[2].append(time.time() - total_start)
                torch.save(
                    model.state_dict(),
                    os.path.join(root_dir, "best_metric_model.pth"),
                )
                print("saved new best metric model")
            print(
                f"current epoch: {epoch + 1} current mean dice: {metric:.4f}"
                f" Fluid-1: {metric_f1:.4f} Fluid-2: {metric_f2:.4f} Fluid-3: {metric_f3:.4f}"
                f"\nbest mean dice: {best_metric:.4f}"
                f" at epoch: {best_metric_epoch}"
            )
    print(f"time consuming of epoch {epoch + 1} is: {(time.time() - epoch_start):.4f}")
total_time = time.time() - total_start

Here is the output from 98th, 99th and 100th epoch

----------
epoch 98/100
1/9, train_loss: 0.9204, step time: 0.6589
2/9, train_loss: 0.9033, step time: 0.6479
3/9, train_loss: 0.9009, step time: 0.6477
4/9, train_loss: 0.9589, step time: 0.6475
5/9, train_loss: 0.9492, step time: 0.6472
6/9, train_loss: 0.9997, step time: 0.6471
7/9, train_loss: 0.9091, step time: 0.6472
8/9, train_loss: 0.9799, step time: 0.6465
9/9, train_loss: 0.8207, step time: 0.6469
10/9, train_loss: 0.6876, step time: 0.3556
epoch 98 average loss: 0.9030
current epoch: 98 current mean dice: 0.1556 Fluid-1: 0.1682 Fluid-2: 0.0054 Fluid-3: 0.1572
best mean dice: 0.1764 at epoch: 63
time consuming of epoch 98 is: 95.6103
----------
epoch 99/100
1/9, train_loss: 0.9993, step time: 0.6640
2/9, train_loss: 1.0000, step time: 0.6486
3/9, train_loss: 0.9885, step time: 0.6474
4/9, train_loss: 1.0000, step time: 0.6466
5/9, train_loss: 0.8556, step time: 0.6462
6/9, train_loss: 0.9631, step time: 0.6467
7/9, train_loss: 0.8788, step time: 0.6479
8/9, train_loss: 0.8513, step time: 0.6463
9/9, train_loss: 0.9607, step time: 0.6468
10/9, train_loss: 0.7564, step time: 0.3552
epoch 99 average loss: 0.9254
current epoch: 99 current mean dice: 0.1612 Fluid-1: 0.1764 Fluid-2: 0.0050 Fluid-3: 0.1587
best mean dice: 0.1764 at epoch: 63
time consuming of epoch 99 is: 93.6048
----------
epoch 100/100
1/9, train_loss: 0.8776, step time: 0.6550
2/9, train_loss: 0.9124, step time: 0.6473
3/9, train_loss: 0.9824, step time: 0.6468
4/9, train_loss: 0.9978, step time: 0.6472
5/9, train_loss: 0.9401, step time: 0.6470
6/9, train_loss: 0.8770, step time: 0.6463
7/9, train_loss: 0.9858, step time: 0.6470
8/9, train_loss: 0.9218, step time: 0.6482
9/9, train_loss: 0.9406, step time: 0.6469
10/9, train_loss: 0.9896, step time: 0.3547
epoch 100 average loss: 0.9425
current epoch: 100 current mean dice: 0.1599 Fluid-1: 0.1735 Fluid-2: 0.0051 Fluid-3: 0.1606
best mean dice: 0.1764 at epoch: 63
time consuming of epoch 100 is: 91.4845

here is my plot of loss and dice score

Any suggestions to reduce loss and improve dice score would be really helpful!

Thank you!

[vikashg]

Hi, Whenever I faced such issues (that the training loss is not reducing at all), it had to do with the data. Can you post a screen shot of your original images. In addition, one thing you should try is to see if the image and labels are aligned after applying the transform chain.

[ankitabuntolia]

My image and label data contain 49 scans per file. so here are the screenshots of scan numbers 2, 12, and 24 of an image with its label.

Thank for the suggestion, I will cross-check after applying the transform chain.

[yiheng-wang-nv]

Hi @ankitabuntolia , you can also try to let the training set be able to overfit by removing all image augmentations (to see if the easiest settings can reach to a very low training loss), for your code the following transforms can be removed:

      RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
      RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
      RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
      Rand3DElasticd(keys=["image", "label"], mode=("bilinear", "nearest"), prob=1.0, sigma_range=(5, 8), magnitude_range=(100, 200),
                    spatial_size=(128, 128, 128), translate_range=(-140, 140), rotate_range=(np.pi / 36, np.pi / 36, np.pi), scale_range=(0.15, 0.15, 0.15), padding_mode="border"),

[ankitabuntolia]

Hi @ankitabuntolia , you can also try to let the training set be able to overfit by removing all image augmentations (to see if the easiest settings can reach to a very low training loss), for your code the following transforms can be removed:

      RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
      RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
      RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
      Rand3DElasticd(keys=["image", "label"], mode=("bilinear", "nearest"), prob=1.0, sigma_range=(5, 8), magnitude_range=(100, 200),
                    spatial_size=(128, 128, 128), translate_range=(-140, 140), rotate_range=(np.pi / 36, np.pi / 36, np.pi), scale_range=(0.15, 0.15, 0.15), padding_mode="border"),

With these changes, it is giving mismatch size of tensor

----------
epoch 1/5
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
[<ipython-input-31-9bcbffb795b7>](https://localhost:8080/#) in <module>()
     25         optimizer.zero_grad()
     26         with torch.cuda.amp.autocast():
---> 27             outputs = model(inputs)
     28             loss = loss_function(outputs, labels)
     29         scaler.scale(loss).backward()

13 frames
[/usr/local/lib/python3.7/dist-packages/monai/networks/layers/simplelayers.py](https://localhost:8080/#) in forward(self, x)
    129 
    130         if self.mode == "cat":
--> 131             return torch.cat([x, y], dim=self.dim)
    132         if self.mode == "add":
    133             return torch.add(x, y)

RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 31 but got size 32 for tensor number 1 in the list.

[ankitabuntolia]

As the error is a mismatch size of the inputs entering the model. So I printed out size of these inputs:

[yiheng-wang-nv]

Hi @ankitabuntolia , I see, the input shape is determined by Rand3DElasticd. Maybe you can use ResizeWithPadOrCropd to achieve the required spatial size for the model. In addition, what is the purpose of using:

      SpatialCropd(keys=['image', 'label'], roi_start=(0,0,0), roi_end=(256,248,49)),
      SpatialCropd(keys=['image_1', 'label_1'], roi_start=(0,248,0), roi_end=(256,496,49),allow_missing_keys=True),
      SpatialCropd(keys=['image_2', 'label_2'], roi_start=(256,0,0), roi_end=(512,248,49),allow_missing_keys=True),
      SpatialCropd(keys=['image_3', 'label_3'], roi_start=(256,248,0), roi_end=(512,496,49),allow_missing_keys=True),
      ConcatItemsd(['image_1', 'image_2', 'image_3'], name='image', dim=1, allow_missing_keys=True),
      DeleteItemsd(['image_1', 'image_2', 'image_3']),

Have you tried to use RandSpatialCropd directly?

[rijobro]

What does this bit do?

      SpatialCropd(keys=['image_1', 'label_1'], roi_start=(0,248,0), roi_end=(256,496,49),allow_missing_keys=True),
      SpatialCropd(keys=['image_2', 'label_2'], roi_start=(256,0,0), roi_end=(512,248,49),allow_missing_keys=True),
      SpatialCropd(keys=['image_3', 'label_3'], roi_start=(256,248,0), roi_end=(512,496,49),allow_missing_keys=True),
      ConcatItemsd(['image_1', 'image_2', 'image_3'], name='image', dim=1, allow_missing_keys=True),
      DeleteItemsd(['image_1', 'image_2', 'image_3']),

If I understand correctly, you're taking 3 sections of your image and joining them together along the x-axis. Does this make sense?