`ViTAutoEnc` hard-codes output patch size to 4

[ArjunNarayanan]

Describe the bug
The ViTAutoEnc constructor hard-codes the output patch size to 4 in this line.

[Edit] I think new_patch_size needs to be equal to patch_size

In particular, this results in an output shape that is different from the input image shape when img_size // patch_size != 4

To Reproduce
Steps to reproduce the behavior:

import torch
from monai.networks.nets import ViTAutoEnc, ViT
net = ViTAutoEnc(
    in_channels=1,
    img_size=64,
    patch_size=8,
)
x = torch.rand([5, 1, 64, 64, 64])
with torch.no_grad():
    out, _ = net(x)

print(out.shape)

Expected behavior
In above example, out.shape = [5,1,128,128,128] when it is expected to be [5,1,64,64,64].

Environment

Ensuring you use the relevant python executable, please paste the output of:

python -c 'import monai; monai.config.print_debug_info()'

================================
Printing MONAI config...                                                                            
================================                                                                    
MONAI version: 1.2.0                                                                                
Numpy version: 1.24.3                                                                               
Pytorch version: 2.0.1+cpu                                                                          
MONAI flags: HAS_EXT = False, USE_COMPILED = False, USE_META_DICT = False                           
MONAI rev id: c33f1ba588ee00229a309000e888f9817b4f1934                                              
MONAI __file__: C:\Users\z003ny8j\Research\Models\ImageFlow\venv\lib\site-packages\monai\__init__.py
                                                                                                    
Optional dependencies:                                                                              
Pytorch Ignite version: NOT INSTALLED or UNKNOWN VERSION.                                           
ITK version: NOT INSTALLED or UNKNOWN VERSION.                                                      
Nibabel version: 5.1.0                                                                              
scikit-image version: NOT INSTALLED or UNKNOWN VERSION.
Pillow version: 9.5.0
Tensorboard version: NOT INSTALLED or UNKNOWN VERSION.
gdown version: NOT INSTALLED or UNKNOWN VERSION.
TorchVision version: NOT INSTALLED or UNKNOWN VERSION.
tqdm version: 4.65.0
lmdb version: NOT INSTALLED or UNKNOWN VERSION.
psutil version: 5.9.5
pandas version: 2.0.2
einops version: 0.6.1
transformers version: NOT INSTALLED or UNKNOWN VERSION.
mlflow version: NOT INSTALLED or UNKNOWN VERSION.
pynrrd version: NOT INSTALLED or UNKNOWN VERSION.

For details about installing the optional dependencies, please visit:
    https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies


================================
Printing system config...
================================
System: Windows
Win32 version: ('10', '10.0.19044', 'SP0', 'Multiprocessor Free')
Win32 edition: Enterprise
Platform: Windows-10-10.0.19044-SP0
Processor: Intel64 Family 6 Model 158 Stepping 13, GenuineIntel
Machine: AMD64
Python version: 3.10.9
Process name: python.exe
Command: ['C:\\Users\\z003ny8j\\AppData\\Local\\anaconda3\\python.exe', '-c', 'import monai; monai.config.print_debug_info()']
Open files: [popenfile(path='C:\\Windows\\System32\\en-US\\kernel32.dll.mui', fd=-1), popenfile(path='C:\\Windows\\System32\\en-US\\KernelBase.dll.mui', fd=-1), popenfile(path='C:\\Windows\\System32\\en-US\\tzres.dll.mui', fd=-1
)]
Num physical CPUs: 6
Num logical CPUs: 12
Num usable CPUs: 12
CPU usage (%): [35.6, 2.4, 7.7, 4.5, 3.2, 2.0, 2.4, 16.1, 16.2, 1.2, 2.4, 58.1]
CPU freq. (MHz): 2592
Load avg. in last 1, 5, 15 mins (%): [0.0, 0.0, 0.0]
Disk usage (%): 80.5
Avg. sensor temp. (Celsius): UNKNOWN for given OS
Total physical memory (GB): 31.7
Available memory (GB): 18.4
Used memory (GB): 13.3

================================
Printing GPU config...
================================
Num GPUs: 0
Has CUDA: False
cuDNN enabled: False

[KumoLiu]

Hi @finalelement, could you please share some comments on this issue?
Thanks in advance!

[wyli]

thanks, it's probably a bug according to Vish, would you like to submit a patch?

[ArjunNarayanan]

Yes, happy to contribute! How should we go about it? We could use an approach similar to what is done in UNETR. For example, if we fix the input patch size to 16, then if we up-sample the ViT features 4 times (since $2^4 = 16$), we will be back in the image dimensions.

[KumoLiu]

Yes, happy to contribute! How should we go about it? We could use an approach similar to what is done in UNETR. For example, if we fix the input patch size to 16, then if we up-sample the ViT features 4 times (since 24=16), we will be back in the image dimensions.

Hi @ArjunNarayanan, since we only have one-time conv here, and two-times deconv here.
We may ensure:

• The size of the input can be divided by patch-size.
• The patch size $\in R^2$.

[finalelement]

@ArjunNarayanan I want to add some highlights here that will be useful for you to formulate the solution.

At the time of development of VitAutoEnc, we were looking to simply add a Decoder to ViT to get the same sized image as output as that of the input. A foremost challenge that we faced was that if we just added a single DeConv there was too much compute overhead as an extremely dense layer gets connected which slows down the network heavily. This led to extremely slow training. Therefore we introduced two layers.

I think that if the parameters are changed for e.g you were experimenting with the patch-size of the tokens, it might make sense to decide the architecture for decoder as well in terms of how many DeConv's for which some experimentation will be required on how many layers, and their hyper-parameters for an optimal training time. Doing this eventually leads to a new architecture.

Adding a DeConv to ViT is an ill-posed problem after all.

Some solutions to consider since the ViTAutoEnc is getting outdated as well
1.) I think a simpler solution here might be to remove the hyper-parameter for the patch-size to avoid confusion for the users.
2.) If we pursue then we need to first make sure that the input is divisible by 4 and that the patch-size is also a factor of 4, after these check are added, we can formulate a relationship for the DeConv's with the patch-size.

Feel free to pursue to the path you like.

[ArjunNarayanan]

Hi @finalelement thanks for sharing these insights, it's very useful. I was wondering if UNETR could be directly used as a 3d image auto-encoder? It also has a ViT as backbone so it might be effectively pre-trained on image tasks?

[finalelement]

Yes you are on point @ArjunNarayanan , the UNETR architecture itself can be pre-trained (using both encoder and decoder of it, that way you should be able to modify token/patch-size as well).

Yes the ViT backbone would get pre-trained in the process, when transferring weights, one would need to be careful that the encoder weights completely get transferred and no layers get skipped out, because if a layer gets skipped the pre-training can actually become detrimental at times.

[wyli]

we merged a fix for this, but please feel free to reopen if there are any issues (e.g. loading existing checkpoints, deconv overhead).