RuntimeError with AttentionUNet

[njneeteson]

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 <module>
    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?

[KumoLiu]

Hi @njneeteson, thanks for reporting, that's a bug in the stride parameter, the current version only supports stride=[2, 2...,
there is a pr to fix it https://github.com//pull/5423
Thanks!

[njneeteson]

Ok thank you! I will hold off on using the attention unet until that patch makes it into a release version.