tensor size mismatched, training custom data with unetr_btcv_segmentation_3d.ipynb

[esther1262]

Hi, My images have dimension 256x256x16, so as the labelmaps. There are 100 images for training and 30 for validation. The labelmaps have 27 labels plus background. I modified some codes in the tutorial to fit my data. I changed some settings in the transformation and UNETR.
After a minute of running, an error popped up that RuntimeError: The size of tensor a (16) must match the size of tensor b (256) at non-singleton dimension 1. How can I fix the issue? The error might come from sliding_window_inference in the validation but I didnt change anything in that function. Calculating epoch loss didnt have issue. I reran the codes with resized option in both train and validation transformation, or just in the train_transforms but issues still occurred.
Thanks!

train_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        EnsureTyped(keys=["image", "label"]),
        NormalizeIntensityd(keys='image'),
        RandBiasFieldd(keys="image", coeff_range=(0.1,0.3), prob=0.5),
        RandRotated(keys=["image", "label"], prob=1, range_z=0.1, mode=['bilinear','nearest']), 
        RandZoomd(keys=["image", "label"], prob=1, min_zoom=0.9, max_zoom=1.1, mode=['area','nearest']),
        Resized(keys=["image", "label"], spatial_size=(256,256,16), mode=['area','nearest'])
    ]
)

val_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        EnsureTyped(keys=["image", "label"]),
        NormalizeIntensityd(keys='image'),
        #Resized(keys=["image", "label"], spatial_size=(256,256,16), mode=['area','nearest'])
    ]
)

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

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

model = UNETR(
    in_channels=1,
    out_channels=28,
    img_size=(256, 256, 16),
    feature_size=32, 
    hidden_size=96,
    mlp_dim=128,
    num_heads=12, 
    pos_embed="perceptron",
    norm_name="instance",
    res_block=True,
    dropout_rate=0.0,
).to(device)

def validation(epoch_iterator_val):
    model.eval()
    with torch.no_grad():
        for batch in epoch_iterator_val:
            val_inputs, val_labels = (batch["image"].to(device), batch["label"].to(device))
            val_outputs = sliding_window_inference(val_inputs,  (96,96,16), 4, model)
            val_labels_list = decollate_batch(val_labels)
            val_labels_convert = [post_label(val_label_tensor) for val_label_tensor in val_labels_list]
            val_outputs_list = decollate_batch(val_outputs)
            val_output_convert = [post_pred(val_pred_tensor) for val_pred_tensor in val_outputs_list]
            dice_metric(y_pred=val_output_convert, y=val_labels_convert)
            epoch_iterator_val.set_description("Validate (%d / %d Steps)" % (global_step, 10.0))
        mean_dice_val = dice_metric.aggregate().item()
        dice_metric.reset()
        
    return mean_dice_val

Errors:

iteration 1/10
Training (2 / 10 Steps) (loss=4.28507):   7%|▋         | 2/27 [01:30<12:30, 30.01s/it]
Validate (X / X Steps) (dice=X.X):   0%|          | 0/28 [00:00<?, ?it/s]
Training (2 / 10 Steps) (loss=4.28507):   7%|▋         | 2/27 [01:30<18:45, 45.03s/it]

Traceback (most recent call last):

  File "/tmp/ipykernel_812587/4198967065.py", line 70, in <module>
    global_step, dice_val_best, global_step_best, resultoutput = train(global_step, train_loader, dice_val_best, global_step_best)

  File "/tmp/ipykernel_812587/4198967065.py", line 46, in train
    dice_val = validation(epoch_iterator_val)

  File "/tmp/ipykernel_812587/4198967065.py", line 8, in validation
    val_outputs = sliding_window_inference(val_inputs, roi_size, sw_batch_size, model)

  File "/home/user/anaconda3/lib/python3.9/site-packages/monai/inferers/utils.py", line 185, in sliding_window_inference
    seg_prob_out = predictor(window_data, *args, **kwargs)  # batched patch segmentation

  File "/home/user/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)

  File "/home/user/anaconda3/lib/python3.9/site-packages/monai/networks/nets/unetr.py", line 195, in forward
    x, hidden_states_out = self.vit(x_in)

  File "/home/user/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)

  File "/home/user/anaconda3/lib/python3.9/site-packages/monai/networks/nets/vit.py", line 110, in forward
    x = self.patch_embedding(x)

  File "/home/user/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)

  File "/home/user/anaconda3/lib/python3.9/site-packages/monai/networks/blocks/patchembedding.py", line 117, in forward
    embeddings = x + self.position_embeddings

RuntimeError: The size of tensor a (16) must match the size of tensor b (256) at non-singleton dimension 1

[KumoLiu]

Hi @esther1262, I think you may have accidentally set the wrong patch_size in sliding_window_inference. (96, 96, 16) -> (256, 256, 16)

val_outputs = sliding_window_inference(val_inputs,  (96,96,16), 4, model)

Hope it helps, thanks!

[esther1262]

Hi @KumoLiu , the program only works if I set the patch_size to (256,256,16), but that's the same as the original image size. Shouldnt the patch size be smaller than the actual image size? That's what's shown in unetr_btcv_segmentation_3d.ipynb. I previously trained a binary segmentation based on the tutorial spleen_segmentation_3d.ipynb. (96,96,16) for the sliding window inference patch size worked well. Thanks.

[esther1262]

Maybe another question, what is the 10 for in epoch_iterator_val.set_description("Validate (%d / %d Steps)" % (global_step, 10.0))? Thx!

[KumoLiu]

Hi @esther1262,

• Yes, patch_size could be smaller than the actual image size, but I think you just set img_size = (256,256,16) which is the input patch size when you initialize the model. In the tutorial, it also set img_size = (96, 96, 96).
• I think 10 is the total iteration number during validation

I hope it helps, thanks!

[esther1262]

Hi @KumoLiu I thought the img_size in UNETR() refers to the dimension of input image, as written in the manual. If I changed that in UNETR() and sliding window inference to (96,96,16). I still got similar error. However, the program worked after adding RandCropByPosNegLabeld in train_transforms. So I am good now. Thanks!