AHNet parameters for 3D CT organ segmentation

[dlabella29]

I am trying to use AHNet in MONAI instead of Unet for 3D segmentation and cannot figure out the parameters to get the model training.

I'm using nifti CT's with nifti kidney segmentations.

https://docs.monai.io/en/latest/networks.html

Unet working...
model = monai.networks.nets.UNet(dimensions=3, in_channels=1, out_channels=2, channels=(32, 64, 128, 256,512),strides=(2, 2, 2, 2), num_res_units=2, norm=Norm.BATCH,dropout=0.1).to(device)

AHNet not working...
model = monai.networks.nets.AHNet(layers=(3,4,6,3),spatial_dims=3,in_channels=1, out_channels=1,psp_block_num=4,upsample_mode='transpose',pretrained=False,progress=True).to(device)

AHNet error log...
File "SegmentHNfinal.py", line 227, in
outputs = model(inputs)
File "/home/user/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
result = self.forward(*input, **kwargs)
File "/home/user/MONAIahnet/MONAI-master/monai/networks/nets/ahnet.py", line 497, in forward
sum0 = self.up0(fm4) + fm3
RuntimeError: The size of tensor a (2) must match the size of tensor b (3) at non-singleton dimension 4

Please let me know if you have any direction on how to get this working.

Thanks!

[Nic-Ma]

Hi @dlabella29 ,

I think the out_channels in your UNet and AHNet configs are different, could you please help double confirm it?
Thanks.

[dlabella29]

Hi @Nic-Ma,

I have tried out_channels values from [1,3] for AHNet, and each one resulted in the same error message.
RuntimeError: The size of tensor a (2) must match the size of tensor b (3) at non-singleton dimension 4

I also tried modifying psp_block_num from [0,4] and still received the same error message.

Would loading the data into cache matter for AHNet? Both ways worked for the UNet. I'm not sure how the rest of the SpleenSegmentation code needs to change when switching from UNet to AHNet on the same dataset.

Is the error possibly being caused by a mismatch here...
"To meet to requirements of the structure, for transpose mode, the input size of the first dim-1 dimensions should be divisible by 2 ** (psp_block_num + 3) and no less than 32. For other modes, the input size of the first dim-1 dimensions should be divisible by 32 and no less than 2 ** (psp_block_num + 3). In addition, at least one dimension should have a no less than 64 size"

Thanks!

[rijobro]

You're using all the default arguments from the constructor, so my guess is the size of the input image is wrong:

import monai
import torch
device='cuda'
model = monai.networks.nets.AHNet().to(device)
model(torch.randn((1, 1, 128, 128, 128)).to(device))

Does the above work?

[dlabella29]

Hi @rijobro,

I ran your attached code without error.

I am also guessing the issue is with the input image.

But, I am not sure where the issue is since UNet is training appropriately without error on the same data set.

It's interesting that UNet is working with the default parameters (and adjustments above), but AHNet is not for the 3Dspleensegmentation file.

Please let me know if I am missing an understanding for the implementation of the two networks.

Thanks!

[rijobro]

Unfortunately, there's no guarantee that if the default parameters of a network are OK with a set of images, then they will also be OK with another default network.

What is the shape fo the images you're trying to use?

[dlabella29]

@rijobro

I have listed 3 different image shapes below from the same training/validation sample dataset...

image shape: torch.Size([512, 424, 87]), label shape: torch.Size([512, 424, 87])
image shape: torch.Size([512, 438, 90]), label shape: torch.Size([512, 438, 90])
image shape: torch.Size([496, 425, 53]), label shape: torch.Size([496, 425, 53])

These were produced 3 times respectively referring to different images/labels by the following...
check_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
check_loader = monai.data.DataLoader(check_ds, batch_size=1)
check_data = monai.utils.misc.first(check_loader)
image, label = (check_data['image'][0][0], check_data['label'][0][0])
print(f"image shape: {image.shape}, label shape: {label.shape}")

These images and their respective shapes worked for UNet. I'm still not sure if any additional formatting or configuration is needed for AHNet.

Thanks again!

[Nic-Ma]

Hi @yiheng-wang-nv ,

Could you please help confirm the data size for AHNet?
I think the dim should be N x 32: https://github.com/Project-MONAI/MONAI/blob/master/monai/networks/nets/ahnet.py#L319

Thanks.

[yiheng-wang-nv]

Hi @dlabella29 , as the link Nic shown, for 3D task, the first 2 dimensions (such as H and W) should be divisible by 2 ** (psp_block_num + 3) and no less than 32, the limitation comes from the PSP module (and the corresponding stride size it used) of the network.