search

Project-MONAI / MONAI

AI Toolkit for Healthcare Imaging
https://monai.io/
Apache License 2.0
5.76k stars 1.06k forks source link

RuntimeError with AttentionUNet #5422

Closed wyli closed 1 year ago

wyli commented 1 year ago

Discussed in https://github.com/Project-MONAI/MONAI/discussions/5421

Originally posted by **njneeteson** October 27, 2022 I'm getting a RuntimeError when using the AttentionUnet in both 2D and 3D. This happens when training using a dataset and training script that works fine with a normal unet, a segan, and a unet-r (in both 2D and 3D) so there must be something specific that I am doing wrong or misunderstanding about how to use the attentionunet properly. Here is a very minimal sample of code to reproduce the problem: ``` import torch from monai.networks.nets.attentionunet import AttentionUnet def test_attention_unet(spatial_dims, channels=(16, 32, 64), img_size=256): attention_unet = AttentionUnet( spatial_dims=spatial_dims, in_channels=1, out_channels=1, channels=channels, strides=[1 for _ in range(len(channels))], dropout=0.0, kernel_size=3, up_kernel_size=3 ) attention_unet.eval() x = torch.zeros(tuple([img_size]*spatial_dims)).unsqueeze(0).unsqueeze(0) print(f"input shape: {x.shape}") with torch.no_grad(): y = attention_unet(x) print(f"output shape: {y.shape}") print("testing in 2D...") test_attention_unet(2) ``` Here is the terminal output: ``` testing in 2D... input shape: torch.Size([1, 1, 256, 256]) Traceback (most recent call last): File "/Users/nathanneeteson/Library/Application Support/JetBrains/PyCharmCE2022.1/scratches/scratch_14.py", line 25, in test_attention_unet(2) File "/Users/nathanneeteson/Library/Application Support/JetBrains/PyCharmCE2022.1/scratches/scratch_14.py", line 20, in test_attention_unet y = attention_unet(x) File "/Users/nathanneeteson/opt/anaconda3/envs/blptl/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl return forward_call(*input, **kwargs) File "/Users/nathanneeteson/opt/anaconda3/envs/blptl/lib/python3.7/site-packages/monai/networks/nets/attentionunet.py", line 256, in forward x_m: torch.Tensor = self.model(x) File "/Users/nathanneeteson/opt/anaconda3/envs/blptl/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl return forward_call(*input, **kwargs) File "/Users/nathanneeteson/opt/anaconda3/envs/blptl/lib/python3.7/site-packages/torch/nn/modules/container.py", line 139, in forward input = module(input) File "/Users/nathanneeteson/opt/anaconda3/envs/blptl/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl return forward_call(*input, **kwargs) File "/Users/nathanneeteson/opt/anaconda3/envs/blptl/lib/python3.7/site-packages/monai/networks/nets/attentionunet.py", line 158, in forward fromlower = self.upconv(self.submodule(x)) File "/Users/nathanneeteson/opt/anaconda3/envs/blptl/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl return forward_call(*input, **kwargs) File "/Users/nathanneeteson/opt/anaconda3/envs/blptl/lib/python3.7/site-packages/torch/nn/modules/container.py", line 139, in forward input = module(input) File "/Users/nathanneeteson/opt/anaconda3/envs/blptl/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl return forward_call(*input, **kwargs) File "/Users/nathanneeteson/opt/anaconda3/envs/blptl/lib/python3.7/site-packages/monai/networks/nets/attentionunet.py", line 159, in forward att = self.attention(g=fromlower, x=x) File "/Users/nathanneeteson/opt/anaconda3/envs/blptl/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl return forward_call(*input, **kwargs) File "/Users/nathanneeteson/opt/anaconda3/envs/blptl/lib/python3.7/site-packages/monai/networks/nets/attentionunet.py", line 139, in forward psi: torch.Tensor = self.relu(g1 + x1) RuntimeError: The size of tensor a (512) must match the size of tensor b (256) at non-singleton dimension 3 ``` I tried to go in and look at the source code to trace the problem but I can't figure out what the problem would be, the code is a bit too complex for me to backtrace to where the issue might be. Does anyone use the AttentionUnet successfully and can see what I am doing wrong? Or is this a bug?
wyli commented 1 year ago

thanks for reporting, that's a bug in the stride parameter, the current version only supports stride=[2, 2..., I submit a patch to fix it https://github.com/Project-MONAI/MONAI/pull/5423