Grad-Cam for nnunetv2

[PARODBE]

Hi,

I'm trying to use monai grad-cam in my nnunetv2 model. First all, I upload my model with the following code:

`import os
import json
import torch
from os.path import join
import nnunetv2
from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import determine_num_input_channels
from nnunetv2.utilities.find_class_by_name import recursive_find_python_class
from nnunetv2.utilities.plans_handling.plans_handler import PlansManager

def load_json_standard(json_path):
with open(json_path, 'r') as f:
return json.load(f)

model_dir = "/home/jovyan/nnunet_data/nnUNet_results/Dataset500_ProstateGland/nnUNetTrainer_250epochs__nnUNetPlans_ResEnc_80G__3d_fullres"
model_path = join(model_dir, "fold_0", "checkpoint_best.pth")

dataset_json = load_json_standard(join(model_dir, 'dataset.json'))
plans = load_json_standard(join(model_dir, 'plans.json'))
plans_manager = PlansManager(plans)

checkpoint = torch.load(model_path, weights_only=False,map_location=torch.device('cpu'))

trainer_name = checkpoint['trainer_name']
configuration_name = checkpoint['init_args']['configuration']
inference_allowed_mirroring_axes = checkpoint.get('inference_allowed_mirroring_axes', None)

parameters = [checkpoint['network_weights']]

configuration_manager = plans_manager.get_configuration(configuration_name)

num_input_channels = determine_num_input_channels(
plans_manager,
configuration_manager,
dataset_json
)

trainer_class = recursive_find_python_class(
join(nnunetv2.path[0], "training", "nnUNetTrainer"),
trainer_name,
'nnunetv2.training.nnUNetTrainer'
)
model = trainer_class.build_network_architecture(
configuration_manager.network_arch_class_name,
configuration_manager.network_arch_init_kwargs,
configuration_manager.network_arch_init_kwargs_req_import,
num_input_channels,
plans_manager.get_label_manager(dataset_json).num_segmentation_heads,
enable_deep_supervision=False
)`

After, as an example I upload data the preprocessed data of the first patient, together along one of the conv layer of the decoder of nnunetv2, and the gradcam module:

`from monai.visualize import GradCAM
from monai.transforms import Resize, NormalizeIntensity, EnsureChannelFirst
from monai.visualize.utils import blend_images
from monai.data import decollate_batch

npz_path = "/home/jovyan/nnunet_data/nnUNet_preprocessed/Dataset500_ProstateGland/nnUNetPlans_3d_fullres/ProstateGland_10000.npz"
data = np.load(npz_path)

image = data['data'] # (C, D, H, W)
label = data['seg'] # (D, H, W) -> Máscara de segmentación

print("Shape del tensor de imagen:", image.shape) # Shape del tensor de imagen: (3, 34, 360, 360)

input_tensor = torch.tensor(image, dtype=torch.float32).unsqueeze(0) # (1, C, D, H, W)

target_layers = ["decoder.stages.5.convs.0.conv"]

cam = GradCAM(nn_module=model, target_layers=target_layers)

cam_map = cam(x=input_tensor, class_idx=1) # Clase 1 = Lesión`

And I obtain the following error:

RuntimeError: The size of tensor a (23) must match the size of tensor b (22) at non-singleton dimension 4

Could you help me, please?

Thanks!