VTUNET_3D causes dimension troubles

[wolke26]

Hi,

I was trying out several models, and it happened with vtunet_3d that I couldnt get to run...

I tried with different input sizes and dummy shapes, but I think the mismatch lies somehwhere further down the line. Unfortunately the original code of VT-UNet is very different that I could not get a quick comparison of the two.

Maybe you have encountered a similar issue in the past? I just want to run their network on some test data.

Thanks in advance!

Below the code:

python train.py

__CUDNN VERSION: 8500
__Number CUDA Devices: 4
__CUDA Device Name: NVIDIA TITAN RTX
__CUDA Device Total Memory [GB]: 25.388515328
Device: cuda
Run: 07_7_2023
Loading configurations from config/aorta/vtunet_3d.yaml
SwinTransformerSys3D expand initial----depths:[2, 2, 2, 1];depths_decoder:[1, 2, 2, 2];drop_path_rate:0.1;num_classes:4;embed_dims:96;window:(7, 7, 7)
---final upsample expand_first---
(178, 64, 64, 40)
current lr: 1e-08
0it [00:00, ?it/s]> /scratch/gwolkerstorf/ASharon_retrain_UNet++/train_for_cluster.py(228)()
0it [00:03, ?it/s]
Traceback (most recent call last):
File "train_for_cluster.py", line 231, in
result = net(img)
File "/home/gwolkerstorf/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
return forward_call(*input, **kwargs)
File "/scratch/gwolkerstorf/ASharon_retrain_UNet++/model/dim3/vtunet.py", line 96, in forward
logits = self.swin_unet(x)
File "/home/gwolkerstorf/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
return forward_call(*input, **kwargs)
File "/scratch/gwolkerstorf/ASharon_retrain_UNet++/model/dim3/vtunet_utils.py", line 2187, in forward
x = self.forward_up_features(x, x_downsample, v_values_1, k_values_1, q_values_1, v_values_2, k_values_2, q_values_2)
File "/scratch/gwolkerstorf/ASharon_retrain_UNet++/model/dim3/vtunet_utils.py", line 1925, in forward_up_features
x = layer_up(x)
File "/home/gwolkerstorf/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
return forward_call(*input, **kwargs)
File "/scratch/gwolkerstorf/ASharon_retrain_UNet++/model/dim3/vtunet_utils.py", line 979, in forward
x = x.view(B, D * 8, H, W, C)
RuntimeError: shape '[1, 32, 4, 4, 1536]' is invalid for input of size 3145728