Home

Awesome

Code for REFERS

This repository provides the official implementation of our paper "Generalized radiograph representation learning via cross-supervision between images and free-text radiology reports", which was published at Nature Machine Intelligence, 2022.

Based on cross-supervised learning (text as supervision), REFERS generates pre-trained models which produce state-of-the-art results on different radiograph benchmarks!

This project is divided into two parts: pre-training and fine-tuning.

Updates

7.7.2022 Model uploaded (ViT-B_16.npz)

7.3.2022 Missing files uploaded: train/val_img_captions_patients.txt (Google Drive)


Installation

1. System requirements

This software was originally designed and run on a system running Ubuntu 18.01, with Python 3.6, PyTorch 1.7, and CUDA 10.1. For a full list of software packages and version numbers, see the Conda environment file environment.yml.

This software leverages graphical processing units (GPUs) to accelerate neural network training and evaluation; systems lacking a suitable GPU will likely take an extremely long time to train or evaluate models. The software was tested with the NVIDIA RTX 2080 TI GPU, though we anticipate that other GPUs will also work, provided that the unit offers sufficient memory.

2. Installation guide

We recommend installation of the required packages using the Conda package manager, available through the Anaconda Python distribution. Anaconda is available free of charge for non-commercial use through Anaconda Inc. After installing Anaconda and cloning this repository, use the conda command to install necessary packages: conda env create -f environment.yml

3. Description of folders


Pre-training

We use MIMIC-CXR-JPG for pre-training. You can acquire more information about this dataset at Johnson et al. MIMIC-CXR-JPG.

1. Setting up datasets

While we provide code to load data for training a deep-learning model, you will first need to download images from the above repositories. We provide a data TXT file, the content inside is a collection of the path name and label of the picture, which will facilitate the reading of the data.

Regarding the downloaded data, I will not introduce too much here, you can go to the corresponding website to view it. Please organize MIMIC dataset in DataProcessed as follows:

./DataProcessed/
    MIMIC/
		train_label.npy
		test_label.npy
		train.txt
		test.txt
		mimic_resize_image.ipynb

After the data are downloaded, we preprocess the data according to our needs to obtain the required image paths and label files (train_label.npy, test_label.npy, train_data.txt, test_data.txt). In order to facilitate the reading of the data, we first preprocess the image using offline scaling and save it locally. You can refer to mimic_resize_image.ipynb for more details.

2. Conduct pre-training using REFERS

This is the command for using 1 GPU for pre-training. It takes about 2 days on one NVIDIA 2080Ti.

python3 pretrain_refers.py --num-gpus-per-machine 1 --cpu-workers 32  --serialization-dir "../data/"  --dist-url "tcp://127.0.0.1:12343"  --checkpoint-every 100

Fine-tuning

1. Datasets

Datasets can be downloaded via following links:

Dataset I COVID-19 Image Data Collection

Dataset II Shenzhen Tuberculosis

Dataset III NIH ChestX-ray

Dataset IV VinBigData Chest X-ray Abnormalities Detection

2. Setting up datasets

Similar to pre-training, you need to first download images from above repositories before running our codes.

Notes: a) Perhaps you need to modify the data paths by yourself. b) The provided training lists are for fine-tuning using 100% labeled data. For fine-tuning with 1% or 10% labeled data, you just to randomly selected 1% or 10% rows (i.e., data) from the provided training lists. We strongly recommend to repeat each experiment for 3 times.

The 4 datasets for fine-tuning in DataProcessed have been organized as follows:

./DataProcessed/
    COVID-19_Image_Data_Collection/
		train.txt
		val.txt
		test.txt
    Shenzhen_Tuberculosis/
		train.txt
		val.txt
		test.txt
    NIH_ChestX-ray/
		train.txt
		val.txt
		test.txt	
    VinBigData_Chest_X-ray/
		train.txt
		val.txt
		test.txt
		data_processed.ipynb

3. Fine-tune models

Data preparation
Training models

Firstly, you should download the pre-trained model via this link and put the model in Fine-tune/checkpoint/.

python3 -m torch.distributed.launch --nproc_per_node=2  --master_addr 127.0.0.2 --master_port 29504  train.py --name caption_100 --stage train --model_type ViT-B_16 --num_classes 1 --pretrained_dir "../checkpoint/refers_checkpoint.pth" --output_dir "./output/" --data_volume '100' --num_steps 500  --eval_batch_size 512 --img_size 224 --learning_rate 3e-2 --warmup_steps 50 --fp16 --fp16_opt_level O2 --train_batch_size 128
python3 -m torch.distributed.launch --nproc_per_node=2  --master_addr 127.0.0.2 --master_port 29504  train.py --name caption_100 --stage train --model_type ViT-B_16 --num_classes 1 --pretrained_dir "../checkpoint/refers_checkpoint.pth" --output_dir "./output/" --data_volume '100' --num_steps 100  --eval_batch_size 512 --img_size 224 --learning_rate 3e-2 --warmup_steps 5 --fp16 --fp16_opt_level O2 --train_batch_size 128
python3 -m torch.distributed.launch --nproc_per_node=2  --master_addr 127.0.0.2 --master_port 29509  train.py --name caption_100 --stage train --model_type ViT-B_16 --num_classes 14 --pretrained_dir "../checkpoint/refers_checkpoint.pth" --output_dir "./output/" --data_volume '100' --num_steps 60000  --eval_batch_size 512 --img_size 224 --learning_rate 1e-2 --warmup_steps 1000 --fp16 --fp16_opt_level O2 --train_batch_size 128
python3 -m torch.distributed.launch --nproc_per_node=2  --master_addr 127.0.0.2 --master_port 29504  train.py  --name caption_100 --stage train --model_type ViT-B_16 --num_classes 14 --pretrained_dir "../checkpoint/refers_checkpoint.pth" --output_dir "./output/" --data_volume '100' --num_steps 100  --eval_batch_size 512 --img_size 224 --learning_rate 3e-2 --warmup_steps 10 --fp16 --fp16_opt_level O2 --train_batch_size 128

3. Evaluation to reproduce the results

Once you have finished the fine-tuning process, you can choose to run python test.pyin corresponding folders to generate the evaluation results (i.e., AUC) on all 4 datasets.

Results should be saved to Results_of_fine-tuning, where we provided our prediction files (for once run) to reproduce simililar ROC curves compared to those provided in the supplementary information. The folders have been organized as follows:

./Results_of_fine-tuning/
    COVID-19_Image_Data_Collection/
		covid_others/
		bacterial_viral/
    Shenzhen_Tuberculosis/
		tuberculosis/
    NIH_ChestX-ray/
		data1/
		data10/
     	data100/
    VinBigData_Chest_X-ray/
		data1/
		data10/
     	data100/

In each sub-folder, we provide the prediction files of 6 approaches (i.e., 5 baselines + ours). Each prediction is organised as True label (1/0) + probability (0-1). We recommend you to use the officially released code of all baselines to generate pre-trained models, on top of which you can conduct fine-tuning on 4 downstream datasets using provided fine-tuning scripts (just remember to modify the codes for loading models).