cvlab-stonybrook/selfmedmae

Code for ISBI 2023 paper "Self Pre-training with Masked Autoencoders for Medical Image Classification and Segmentation"

SelfMedMAE: Self Pre-training with Masked Autoencoders for Medical Image Analysis

Necessary Packages and Versions

pytorch=1.7.1=py3.8_cuda10.1.243_cudnn7.6.3_0
torchvision=0.8.2=py38_cu101
monai-weekly==0.9.dev2152
nibabel==3.2.1
omegaconf==2.1.1
timm==0.4.12

Preparation

1. Install PyTorch, timm and MONAI.
2. Download the BTCV and MSD_BraTS data.
3. Install Wandb for logging and visualizations.

Stage 1: MAE Pre-Training

The run scripts are in directory scripts

python main.py \
        configs/mae3d_btcv_1gpu.yaml \
        --mask_ratio=0.125 \
        --run_name='mae3d_sincos_vit_base_btcv_mr125'

The default configurations are set in configs/mae3d_btcv_1gpu.yaml. You can overwrite the configurations by passing arguments with the corresponding key names through the command line, e.g., mask_ratio. We use Wandb to monitor the training process and visualize the masked reconstruction. During the training, the output including checkpoints and Wandb local files are all stored in the specified output_dir value in the configurations. The core MAE codes locate in lib/models/mae3d.py.

Stage 2: UNETR Fine-tuning

The run scripts is in directory scripts

python main.py \
        configs/unetr_btcv_1gpu.yaml \
        --lr=3.44e-2 \
        --batch_size=6 \
        --run_name=unetr3d_vit_base_btcv_lr3.44e-2_mr125_10ke_pretrain_5000e \
        --pretrain=$YOUR Pre-Trained MAE Checkpoint$

The core UNETR codes locate in lib/models/unetr3d.py.

selfmedmae's Issues

Did you calculate the MSE or FourierMSE Loss on all tokens or only on masked tokens?

Hello, your work on mae is great! Could you please give more details on the MSE loss calculation?
In "lib/models/mae.py", it seems there are many types of loss while some of them are annotated. Which is the one adopted?

Very Dark Input?

Hi,

Quick question - when looking at the visualizations generated, I notice that sometimes the input is very dark / black. If you zoom in closely, it's apparent that the underlying data (BTCV) is the same. Did this happen to you and is there an explanation for this? In my case, this occurred on the first epoch. This also happens over time once the loss has been NaN for several tens or hundreds of epochs.

Any help is appreciated!

Version

Hi

could you please mention the version of the libraries you used with complete instruction of data preprocessing and etc. I am getting lot of errors like "No module named 'timm.models.layers.helpers'". Installed timm but still getting errors. Please make a requirements.txt/yml

regards
Nazib

A issue while training on my own liver tumor dataset.

Hello，developers of SelfMedMAE.
Your work is valuable enlightening me so much and I am trying to make some improvements to improve the effect on my liver tumor dataset.
Considering that my dataset has only 2 kinds of classes but BTCV has 14 kinds, I has made some changes to the code for the 2 output channels.And the training epoch is set to 1000 because 200 data will cost a lot of time. However, the test Dice was just under 0.6, whereas I was able to achieve around 0.66 when training on top of a normal UNETR without any pre-train.
Now I am trying to do some pre-train and hoping to see some changes. If the accuracy increases, I can keep do other imporvment on the SelfMedMAE.
Thank you very much!

cxr_train_Code

Hello, could you please supplement the code for the cxr dataset?Thank you very much

NaN Loss

Hi,

Thanks for releasing your code publicly. I have been encountering problems training this using the BTCV dataset as well as my own during training. I notice that the loss quickly becomes NaN before 100 epochs. Did you encounter this same issue? I am using the same library packages, pytorch and python versions. If you did encounter this during your experiments, will you please share your approach to handling it?

Any help is greatly appreciated!

There doesn't has BTCV\transunet.json

I run with BTCV, but there is a error,

ValueError: Data list file data\BTCV\transunet.json does not exist.

What can I do? Where can download the transunet.json.

About maevitbackbone and unetr decoder

Thank you for your paper,

1. When you apply the MAE vit, in MAE architecture there are encoder and decoder, did you used mae encoder ? did you use mask ratio of 0 to make same size of token embeddings to reconstruct to original size.
   if we have mask ratio, then token embeding size get smaller as you know.

2. did you freeze mae encoder which consist patch, block, normalization.

3. Finally, the architecture will be Mae encoder (before merging mask tokens in mae decoder, not using this) + Unetr decoder is the architecture that you used.?

thanks

library versions

could anyone who share their version of library in this project. Much appraciate

exact settings of BTCV dataset in pretraining

I wonder the warm-up epochs of btcv dataset. The paper does not mention this, and the configs is weird. The total epochs is 800 in the paper, but the configs there is 10000. I'm not sure the warm-up=100 (in the configs) is right for btcv dataset.

A issus of MAEViTEncoder's code

Hi!Developers of SlefMedMAE!
While I checking the code of MAEViTEncoder, I see self.patch_embed = embed_layer(img_size=patch_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim). I'd like to know why img_size uses patch_size instead of the size of the img?