apple1986 / levit-unet Goto Github PK

您好，我在运行LeViT-UNet模型时遇到了一下错误：
Traceback (most recent call last):
File "D:/Program Files/JetBrains/PyCharm Community Edition 2023.1.4/plugins/python-ce/helpers/pydev/pydevd.py", line 1496, in _exec
pydev_imports.execfile(file, globals, locals) # execute the script
File "D:\Program Files\JetBrains\PyCharm Community Edition 2023.1.4\plugins\python-ce\helpers\pydev_pydev_imps_pydev_execfile.py", line 18, in execfile
exec(compile(contents+"\n", file, 'exec'), glob, loc)
File "E:\test\test2_muscle and adipose segmentation\test2_SL\train.py", line 364, in
Train(train_root)
File "E:\test\test2_muscle and adipose segmentation\test2_SL\train.py", line 211, in Train
outputs = net(inputs)
File "D:\ProgramData\anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "E:\test\test2_muscle and adipose segmentation\test2_SL\model_LeViT_UNet_128s.py", line 583, in forward
x = self.block_1(x) # torch.Size([1, 196, 128])--> torch.Size([1, 196, 128]) : CNN->Trans
File "D:\ProgramData\anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "D:\ProgramData\anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\container.py", line 204, in forward
input = module(input)
File "D:\ProgramData\anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "E:\test\test2_muscle and adipose segmentation\test2_SL\model_LeViT_UNet_128s.py", line 151, in forward
return x + self.m(x)
File "D:\ProgramData\anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "E:\test\test2_muscle and adipose segmentation\test2_SL\model_LeViT_UNet_128s.py", line 216, in forward
(q @ k.transpose(-2, -1)) * self.scale
RuntimeError: The size of tensor a (1024) must match the size of tensor b (196) at non-singleton dimension 3

LeViT-UNet网络调用命令为：net = Build_LeViT_UNet_128s(num_classes=5)
我的输入数据的维度是(4，1，512，512)，num_classes是5，请问这个错误是由什么引起的呢？

Is there a version for 3D data?