/Combining-chest-X-rays-and-EHR-data-ARF

Primary LanguagePython

Overview

This is the code repository for the manuscript "Combining chest X-rays and EHR data using machine learning to diagnose acute respiratory failure".

Directory structures

  • dataset/ contains data loaders

  • model/ contains model loaders

  • checkpoint/ is where model checkpoints are saved

Download and Process the Data

Follow directions to download the MIMIC-IV, MIMIC-CXR, and CheXpert datasets

Note: code to extract an AHRF cohort from MIMIC will be available soon.

Running the code

Config file

Pre-specified arguments can be set in config.json:

Required arguments:

  • csv_file: Path to metadata file.
  • checkpoint: Path to file location where model checkpoints will be saved.
  • labels: column name in the metadata files of the classes. These can be separated by "|" (e.g., CHF|Pneumonia|COPD)
  • rotate degrees: degrees of rotation to use for random rotation in image augmentation.
  • disk: if disk = 1, all images will be loaded into memory before training. Otherwise during training images will be fetched from disk.
  • mask : mask = 1 if masked loss will be used (i.e., if there are missing labels). All missing labels in the metadata file should be set to -1.
  • early_stop: early_stop = 1 if early stopping criteria will be used. Otherwise model will train to 3 epochs.
  • pretrain: Whether or not to use an initialization. If pretrain is "yes", then ImageNet initialization will be used unless a pretrain file is specified. Otherwise, pretrain should be "random"
  • pretrain_file: file path to pretrained model (i.e., pretrained model on MIMIC-CXR and CheXpert)
  • pretrain_classes: number of labels pretrain model had
  • freeze_all: 1 or 0: whether or not to freeze all the layers but the classifier in the DenseNet
  • loader_names : list of split names (i.e., ["train", "valid", "test"]). You do not have to include "test".

Training a model

The following exmple code will train a model using train.py. Each run requires that a model_name and model_type be specificied. There are pre-specified in the config file along with other parameters (described in further detail below). Models will be saved in the directory chexpoint/model_type/model_name.

Other non-required arguments are:

Arguments

  • gpu: specify the gpu numbers to train on, default is 0 and 1.
  • budget: number of hyperparameter combinations to try. Default is 50.
  • repeats: number of seed initializations to try. Default is 3.
  • save_every: for pretraining on MIMIC-CXR and CheXpert. Number of iterations to complete before saving a checkpoint. Default is None and will save after every epoch.
  • save_best_num, for pretraining on MIMIC-CXR and CheXpert. Number of top checkpoints to save (based on best AUROC performance on the validation set). Default is 1.
  • optimizer: optimzier to use. Default is "sgd", but can also choose "adam" for pretraining on MIMIC-CXR and CheXpert.
python train.py --model_type example_model_type --model_name example_model_name

Pretraining

To train a model on MIMIC-CXR and CheXpert, you'll want to use the save_every, save_best_num, and optimizer arguments. This will train on an ImageNet initialized model:

python train.py --model_type example_model_type --model_name example_model_name --save_every 4800 --save_best_num 10 --optimizer adam

Fine-tuning a model on the AHRF cohort:

To train a DenseNet model after pretraining on either MIMIC-CXR/CheXpert, you'll need to specify the file location of the pretrained model in the config file, as well as the number of classes in the pretrained model.