Issue with ROCAUCMetric in MONAI: Prediction Shape Error

[jsmorac]

I am encountering an issue when using the ROCAUCMetric from MONAI to evaluate my segmentation model's performance. When I run my script, I get the following error:
line 135, in metrics_labels
auc_value = auc_metric.aggregate().
ValueError: Predictions should be of shape (batch_size, num_classes) or (batch_size, ), got torch.Size([1, 11, 240, 185, 264]).

Here is the code snippet where the ROCAUCMetric is set up and used:
auc_metric = ROCAUCMetric(average="none")

post_label = AsDiscrete(to_onehot=11)
post_pred = AsDiscrete(argmax=True, to_onehot=11)

post_prob = Compose([Activations(softmax=True)])
y_trans = Compose([AsDiscrete(to_onehot=11)])
devices = torch.device("cpu")

model = UNETR(
in_channels=1,
out_channels=11,
img_size=(96, 96, 96),
feature_size=16,
hidden_size=768,
mlp_dim=3072,
num_heads=12,
pos_embed="perceptron",
norm_name="instance",
res_block=True,
dropout_rate=0.0,
).to(devices)
model.eval()

with torch.no_grad():
for val_data in val_loader:
val_inputs, val_labels = (val_data["image"].cpu(), val_data["label"].cpu())
val_outputs = sliding_window_inference(val_inputs, (96, 96, 96), 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_labels_prob = [y_trans(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]
    val_output_prob = [post_prob(val_pred_tensor) for val_pred_tensor in val_outputs_list]
    
    auc_metric(y_pred=val_output_prob, y=val_labels_prob)

auc_value = auc_metric.aggregate()
auc_metric.reset()
print(auc_value)

How can I resolve the issue related to the prediction shape to make it compatible with the ROCAUCMetric?
Are there any specific guides or examples on how to prepare the predictions for ROC AUC calculation in a segmentation context in MONAI?
Do I need to adjust the shape of the model outputs or perform any additional transformations before passing the data to the ROCAUCMetric?
Any help or suggestions would be greatly appreciated. Thank you!

[KumoLiu]

Hi @jsmorac,

Predictions should be of shape (batch_size, num_classes) or (batch_size, ), got torch.Size([1, 11, 240, 185, 264]).

From the error message, you can that the prediction should be shape of (batch_size, num_classes), which is a typical classification model output instead of segmentation model.

Thanks.

[jsmorac]

Thank you for the response. I understand that the predictions should be in the shape (batch_size, num_classes) instead of (batch_size, 11, 240, 185, 264). My question is how can I transform my predictions to have the correct shape. Should I flatten the dimensions or perform any additional transformations? Any guidance or examples would be very helpful. Thank you.