The code doesn't work when H×W of input is different than 256×256.

[shahzad-ali]

Following is the code for reproducibility:

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model = UNet_icl(in_chns=3, class_num=7).to(device)

x_lab = torch.tensor(np.random.rand(4, 3, 512, 512), dtype=torch.float32, device=device)
x_unlab = torch.tensor(np.random.rand(4, 3, 512, 512), dtype=torch.float32, device=device)
p = model(x_lab, x_unlab, inference=False)

I am trying to train the model for the input of shape 4×3×512×512, but it fails 😔, and the following error is shown:

Traceback (most recent call last):
File "C:\Users\User\PycharmProjects\sarcopenia\mva\icl_unet.py", line 681, in
p = model(x_lab, x_unlab, inference=False)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\User\PycharmProjects\sarcopenia\venv_3.11\Lib\site-packages\torch\nn\modules\module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\User\PycharmProjects\sarcopenia\venv_3.11\Lib\site-packages\torch\nn\modules\module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\User\PycharmProjects\sarcopenia\mva\icl_unet.py", line 448, in forward
feat_Maps_lab, updated_Qs_lab = self.sspa(feats_lab, 'labeled')
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\User\PycharmProjects\sarcopenia\venv_3.11\Lib\site-packages\torch\nn\modules\module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\User\PycharmProjects\sarcopenia\venv_3.11\Lib\site-packages\torch\nn\modules\module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\User\PycharmProjects\sarcopenia\mva\icl_unet.py", line 517, in forward
updated_guided_Q, attn_map = self.class_decoders[i_layer](next_guided_Q, tok_feats)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\User\PycharmProjects\sarcopenia\venv_3.11\Lib\site-packages\torch\nn\modules\module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\User\PycharmProjects\sarcopenia\venv_3.11\Lib\site-packages\torch\nn\modules\module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\User\PycharmProjects\sarcopenia\mva\icl_unet.py", line 571, in forward
attn = attn + self.drop_path(self.mlp2(self.norm3(attn)))
^^^^^^^^^^^^^^^^
File "C:\Users\User\PycharmProjects\sarcopenia\venv_3.11\Lib\site-packages\torch\nn\modules\module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\User\PycharmProjects\sarcopenia\venv_3.11\Lib\site-packages\torch\nn\modules\module.py", line 1527, in _call_impl
return forward_call(args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\User\PycharmProjects\sarcopenia\venv_3.11\Lib\site-packages\torch\nn\modules\normalization.py", line 196, in forward
return F.layer_norm(
^^^^^^^^^^^^^
File "C:\Users\User\PycharmProjects\sarcopenia\venv_3.11\Lib\site-packages\torch\nn\functional.py", line 2543, in layer_norm
return torch.layer_norm(input, normalized_shape, weight, bias, eps, torch.backends.cudnn.enabled)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Given normalized_shape=[1024], expected input with shape [, 1024], but got input of size[4, 7, 8, 4096]

I have tried for several heights and widths and can confirm that it only works for 256×256 😢.

Your help in making the shared code work for 5125×512 will be greatly appreciated 🙏🏼

[zhuye98]

Hi, Shahzad. The normalized_shape in layer normalization needed to be adjusted to accommodate the change in input dimensions. You should update the param in UNet_icl from input_resolution: [16, 32, 64, 128, 256] to input_resolution: [32, 64, 128, 256, 512]. However, training images at this higher resolution may require an NVIDIA Tesla A100 80G. I hope this will help.