AssertionError: ground truth has different shape (torch.Size([1, 3, 128, 128, 127])) from input (torch.Size([1, 3, 128, 128, 128]))

[ankitabuntolia]

Problem:

After running 122 epoch, at 123rd epoch DiceLoss throwing an error, AssertionError: ground truth has different shape (torch.Size([1, 3, 128, 128, 127])) from input (torch.Size([1, 3, 128, 128, 128]))

My question is why I am getting this error after running 122 epochs successfully? and How can I get rid of it?

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]).

My Transform

train_transform = Compose(
    [
      LoadImaged(keys=["image", "label"]),
      EnsureChannelFirstd(keys="image"),
      ConvertToMultiChannelBasedOnBratsClassesd(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, 2.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="image"),
      NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
      # SpatialPadd(keys=['image', 'label'], spatial_size= [512, 512, 49]),# it will result in 512x512x49
      EnsureTyped(keys=["image", "label"]),
    ]
)
val_transform = Compose(
    [
      LoadImaged(keys=["image", "label"]),
      EnsureChannelFirstd(keys="image"),
      ConvertToMultiChannelBasedOnBratsClassesd(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, 2.0), mode=("bilinear", "nearest")),
      CropForegroundd(keys=["image", "label"], source_key="image"),
      NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
      # SpatialPadd(keys=['image', 'label'], spatial_size= [512, 512, 49]),# it will result in 512x512x49
      EnsureTyped(keys=["image", "label"]),
    ]
)

I am following Brain tumor 3D segmentation with MONAI but instead of using its model, I am using it from Tutorial: Spleen 3D segmentation with MONAI with adjustment provided by #415.

Here is my model:

max_epochs = 600
val_interval = 1
VAL_AMP = True

# standard PyTorch program style: create Unet Model, DiceLoss and Adam optimizer
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-4, 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

I am using training code from Brain tumor 3D segmentation with MONAI

best_metric = -1
best_metric_epoch = -1
best_metrics_epochs_and_time = [[], [], []]
epoch_loss_values = []
metric_values = []
metric_values_tc = []
metric_values_wt = []
metric_values_et = []

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_tc.append(metric_f1)
            metric_f2 = metric_batch[1].item()
            metric_values_wt.append(metric_f2)
            metric_f3 = metric_batch[2].item()
            metric_values_et.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

At 123rd epoch, my training is throwing this error:

----------
epoch 123/600
/usr/local/lib/python3.7/dist-packages/torch/functional.py:445: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  ../aten/src/ATen/native/TensorShape.cpp:2157.)
  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]
/usr/local/lib/python3.7/dist-packages/torch/functional.py:445: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  ../aten/src/ATen/native/TensorShape.cpp:2157.)
  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]
/usr/local/lib/python3.7/dist-packages/torch/functional.py:445: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  ../aten/src/ATen/native/TensorShape.cpp:2157.)
  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]
/usr/local/lib/python3.7/dist-packages/torch/functional.py:445: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  ../aten/src/ATen/native/TensorShape.cpp:2157.)
  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]
1/9, train_loss: 0.9467, step time: 0.7273
2/9, train_loss: 0.8993, step time: 0.6836
3/9, train_loss: 0.8767, step time: 0.7431
4/9, train_loss: 0.9954, step time: 0.7097
5/9, train_loss: 1.0000, step time: 0.6605
6/9, train_loss: 1.0000, step time: 0.6821
7/9, train_loss: 0.9506, step time: 0.6605
8/9, train_loss: 0.9365, step time: 0.6544
9/9, train_loss: 0.9544, step time: 0.6499
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
[<ipython-input-55-2ad040568dd7>](https://localhost:8080/#) in <module>()
     26         with torch.cuda.amp.autocast():
     27             outputs = model(inputs)
---> 28             loss = loss_function(outputs, labels)
     29         scaler.scale(loss).backward()
     30         scaler.step(optimizer)

1 frames
[/usr/local/lib/python3.7/dist-packages/monai/losses/dice.py](https://localhost:8080/#) in forward(self, input, target)
    155 
    156         if target.shape != input.shape:
--> 157             raise AssertionError(f"ground truth has different shape ({target.shape}) from input ({input.shape})")
    158 
    159         # reducing only spatial dimensions (not batch nor channels)

AssertionError: ground truth has different shape (torch.Size([1, 3, 128, 128, 127])) from input (torch.Size([1, 3, 128, 128, 128]))

[wyli]

I think it makes sense that the preprocessing defined here is unstable.

Because of the CropForegroundd transform, the spatial shape becomes dependent on the intensity values (bounding box of positive as foreground), and because of the random deformation transform, the intensities will have some randomness. As a result, the preprocessing output could be pairs of spatial shape (128,128,127) image and label. Note that the unet used here with input spatial size (128, 128, 127) will output a spatial (128,128,128) tensor, and eventually the error is at computing loss function.

[ankitabuntolia]

Ok, I will try it again without CropForegroundd , it will require a couple of hours to see if the error pops up again.

I have additional some questions:

1. Is there any best practice/suggestion about using various transforms? Like which one to use when and in which order? I found one tutorial 3D-Augmentation but it does not include all offered transform by Monai.

2. SpatialCropd just gives one cropped part from an image but does not create multiple clipping. I would like to create 4 patches from my single input image. Is there any transform with clipping functionality?

3. As per another discussion (Model loss decreases but validation DICE is always 0 #1727), it has been suggested to use DiceMetric rather than DiceLoss but in that case with multi-label segmentation, include_background should be include_background = False?

4. In my previous training, it ran for 122 epochs (at 123rd epoch I got an error) but the learning is almost like nothing. For this training run, parameters and transforms are available at the main body of this query.

Thank you Wyli!

[ishraq-dagamseh]

@ankitabuntolia
try this code for post process

from torch import nn
predicted_masks = nn.functional.interpolate(predicted_masks,
size=( 128, 128),# decide the input size
mode='bilinear',
align_corners=False)