[Shape mismatch] model prediction changing dimension

[aalhayali]

Hi!

I'm currently training a 3D-UNet model. The image and label size I'm working with are of sagittal MRI (512,512,7). I'm using the following transforms for both train and test (minus the augmentations of course):

train_transform = Compose([
    LoadImaged(keys = ["image", "label"]),
    EnsureChannelFirstd(keys = ["image", "label"]),
    Spacingd(keys=["image","label"], pixdim=(2,1.5,1.5)),
    ScaleIntensityRanged(keys="image", a_min =265., a_max =1555. , b_min = 0.0, b_max = 1.0, clip=True), #apply to image only not the label
    CropForegroundd(keys=["image","label"], source_key="image"),
    Resized(keys=["image","label"], spatial_size=[7,256,256]),
    #Augmentations
    RandAffined(keys=['image', 'label'], prob=0.5, translate_range=10), 
    RandRotated(keys=['image', 'label'], prob=0.5, range_x=10.0),
    RandGaussianNoised(keys='image', prob=0.5),
    ToTensord(keys=["image", "label"])
])

The output image size after loading the dataset to a dataloader and using a batch size of 2 is (2, 1, 7, 512, 512). I I understand that the 1 is a placeholder (though don't fully understand it's purpose, I'd appreciate if someone can clarify it). When passing the dataloaders into the training loop, the image is increasing by 1 in the second dimension (index 1). Here's a snippet of the training code:

    best_metric = -1
    best_metric_epoch = -1
    save_loss_train = []
    save_loss_test = []
    save_metric_train = []
    save_metric_test = []

    for epoch in range(max_epochs):
        print("-" * 10)
        print(f"epoch {epoch + 1}/{max_epochs}")
        model.train()
        train_epoch_loss = 0
        train_step = 0
        epoch_metric_train = 0
        for batch_data in train_loader:
            
            train_step += 1

            volume = batch_data["image"]
            print(f"image shape:{volume.shape}")
            label = batch_data["label"]
            print(f"label shape: {label.shape}")
            label = label != 0
            volume, label = (volume.to(device), label.to(device))

            outputs = model(volume)
            
            train_loss = loss(outputs, label)
            optim.zero_grad()
            train_loss.backward()
            optim.step()

            train_epoch_loss += train_loss.item()
            print(
                f"{train_step}/{len(train_loader) // train_loader.batch_size}, "
                f"Train_loss: {train_loss.item():.4f}")
            
            print(f"outputs shape is: {outputs.shape} \nlabels shape is:{label.shape}")

            train_metric = dice_metric(outputs, label)
            epoch_metric_train += train_metric
            print(f'Train_dice: {train_metric:.4f}')

        print('-'*20)

Image and label shape before passing volume to model were (2, 1, 7, 256, 256), however, output generates a prediction with the following shape: (2, 2, 7, 256, 256), while the label's shape remains as it was originally. This of course is generating a ValueError: y_pred and y should have same shapes, got (2, 2, 7, 256, 256) and (2, 1, 7, 256, 256). due to the mismatch, precisely at the train_loss variable. May I know why I'm getting this dimensionality increase? I looked around for a while and have yet to see this issue.

Thanks!

[KumoLiu]

Hi @aalhayali, 1 in (2, 1, 7, 512, 512) is the channel dim. For all the data we put into the transform we need to make sure that its shape is (C, H, W, [D]). After Dataloader, it will become (B, C, H, W, [D]), so in your case, it seems your model needs a one-channel input and 2 channel output. It may be the reason why 1 becomes 2 in the channel dim. You can check the initialization of your model, if you only want one channel output, just set the output channel to 1, then the output shape will be (2, 1, 7, 256, 256).

Hope it can help you, thanks!

[aalhayali]

Thanks for getting back!

Understood, the 1 in the dimension inserts the channel number which I realised lacks in medical images (compared to regular vision tasks). Since it's a semantic segmentation task, the output should be 2 (foreground and background). Here's the code I have (which has out_channels == 2):

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=2,
    channels=(4, 8, 16, 32, 64),
    strides=(1, 1, 1, 1),
    num_res_units=2,
    norm = Norm.BATCH
  ).to(device)

loss_function = DiceCELoss(to_onehot_y=True, softmax=True)
torch.backends.cudnn.benchmark = True
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5)
dice_metric = DiceMetric(include_background=True, reduction="mean")
scaler = torch.cuda.amp.GradScaler()
dice_metric_batch = DiceMetric(include_background=True, reduction="mean_batch")

I changed out_channels to 1 and training is working now (still in epoch 1/20. Although it doesn't make sense to me why out_channels == 1 would work, as this means the output would be one class, so how will a mask be generated? I may be misunderstanding the inner workings of it, so I'd appreciate it if you could clarify it!

[aalhayali]

Follow-up:

When setting the out_channels = 1, the training loop ran, however, the dice didn't increase (in fact, avg. is 0.05) and loss didn't decrease (0.8). My model therefore isn't learning. Any suggestions on why that is? My training and testing loops are fine, so the only thing that's incorrect in this case is the model itself. Attached are the curves:

Could it be the resized transformation mentioned earlier Resized(keys=["image","label"], spatial_size=[7,256,256])? I've read some tutorials that keep it in the form of (H, W, Slices), however, when I do that, the images are essentially corrupt. For example, when plotting the 3rd slice (index = 2 in original, 4 in transformed), the result is the following:

However, when I keep the transformation in the way I had it earlier (Slices, H, W), the output images are just fine:

[KumoLiu]

Hi @aalhayali, if you only have background and foreground, then two channels means your labels are one-hot formatted and one channel means your labels aren't one-hot.
For resize part, you must set spatial_size to be consistent with your input shape.
Hope it can help you, thanks!