PRETRAIN.md

Pre-training

Fully-supervised with ImageNet (CNNs & ViTs)

Please refer to TorchVision Offical for CNN models and DeiT: Data-efficient Image Transformers for ViTs. The used checkpoints are provided in this Table.

Self-supervised with ImageNet (CNNs & ViTs)

For the CNN-based methods (MoCo V2, BYOL and SwAV), please refer to the official repositories MoCo Official, SwAV Official, and BYOL Official for pre-training details. The used checkpoints are provided in this Table.

For MAE pretraining for ViT on ImageNet, please refer to the MAE Official.

Self-supervised with Chest X-rays (CNNs)

For MoCo V2, we use the code from MoCo Official and just replace the dataloader with the one provided in this repo.

For MAE, we implement a version based on the U-Net framework using Segmentation models with pretrained backbones. PyTorch. The pretraining task is the reconstruction task similar to the original MAE.

The pretraining command on a 8-GPU machine is shown as below:

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 OMP_NUM_THREADS=1 \
python -m torch.distributed.launch --nproc_per_node=8 \
 --use_env main_pretrain_multi_datasets_xray_cnn.py \
 --output_dir ${SAVE_DIR} \
 --log_dir ${SAVE_DIR} \
 --batch_size 256 \
 --model 'densenet121' \
 --mask_ratio 0.75 \
 --epochs 800 \
 --warmup_epochs 40 \
 --blr 1.5e-4 --weight_decay 0.05 \
 --num_workers 8 \
 --input_size 224 \
 --random_resize_range 0.5 1.0 \
 --datasets_names chexpert chestxray_nih

Self-supervised with Chest X-rays (ViTs)

We pretrain ViTs with MAE following the official repo but with **a customized recipe **(please refer to the paper for more details). Two sample commands are provided below.

To pretrain ViT-S on CheXpert and NIH Chest-Xray:

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 OMP_NUM_THREADS=1 \
python -m torch.distributed.launch --nproc_per_node=8 \
 --use_env main_pretrain_multi_datasets_xray.py \
 --output_dir ${SAVE_DIR} \
 --log_dir ${SAVE_DIR} \
 --batch_size 256 \
 --model mae_vit_small_patch16_dec512d2b \
 --mask_ratio 0.90 \
 --epochs 800 \
 --warmup_epochs 40 \
 --blr 1.5e-4 --weight_decay 0.05 \
 --num_workers 8 \
 --input_size 224 \
 --mask_strategy 'random' \
 --random_resize_range 0.5 1.0 \
 --datasets_names chexpert chestxray_nih

To pretrain ViT-B on CheXpert, NIH Chest-Xray and MIMIC:

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 OMP_NUM_THREADS=1 \
python -m torch.distributed.launch --nproc_per_node=8 \
 --use_env main_pretrain_multi_datasets_xray.py \
 --output_dir ${SAVE_DIR} \
 --log_dir ${SAVE_DIR} \
 --batch_size 256 \
 --model mae_vit_base_patch16_dec512d8b \
 --mask_ratio 0.90 \
 --epochs 800 \
 --warmup_epochs 40 \
 --blr 1.5e-4 --weight_decay 0.05 \
 --num_workers 8 \
 --input_size 224 \
 --mask_strategy 'random' \
 --random_resize_range 0.5 1.0 \
 --datasets_names chexpert chestxray_nih mimic_cxr