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