updated readme with info for training

model_zoo/DROID-RV/readme.md

@@ -1,4 +1,6 @@
-# DROID-RV Inference Example
+# DROID-RV Overview
+
+## Inference Example
 
 This is a simple example script demonstrating how to load and run the DROID-RV and DROID-RVEF models. Model training and inference was performed using the code provided in the ML4H [model zoo](https://github.com/broadinstitute/ml4h/tree/master/model_zoo/DROID). The example below was adapted from the DROID inference code.
 
@@ -43,3 +45,47 @@ Model outputs for DROID-RVEF take the form:
     [["Female", "Male"]]
 ]
 ```
+
+## Training Example
+
+Data preprocessing and model training was performed using the DROID training recipe as seen in the [DROID model zoo entry](https://github.com/broadinstitute/ml4h/tree/master/model_zoo/DROID). The Docker image described above can also be used for model training.
+
+Requirements:
+- Wide file: Parquet file with one row per TTE video, and columns corresponding to sample identifier, patient split, and outcome(s) of interest
+- Splits file: JSON file with keys corresponding to "patient_train", "patient_valid", "patient_test" and values corresponding to lists of patient identifiers
+- LMDB folder: prepared as described in the [DROID model zoo entry](https://github.com/broadinstitute/ml4h/tree/master/model_zoo/DROID)
+- Movinet checkpoint: movinet_a2_base in the [DROID-RV model zoo entry](https://github.com/broadinstitute/ml4h/tree/master/model_zoo/DROID-RV)
+- (Optional) pretrained checkpoint: for example https://github.com/broadinstitute/ml4h/tree/master/model_zoo/DROID-RV/droid_rv_checkpoint
+
+For valid arguments for view, Doppler, quality, and canonical axis, see model_zoo/DROID/echo_defines.py.
+
+The command below is an example of how the training script can be used:
+
+```
+python model_zoo/DROID/echo_supervised_training_recipe.py \
+    --n_input_frames 16 \
+    --output_labels age \ # output labels must be present as a column in the wide file
+    --output_labels sex \ 
+    --output_labels rvedd \
+    --output_labels rv_size \
+    --output_labels rv_function \
+    --output_labels_types rcrcc \ # r = regression or c = classification; must be given in same the order as output labels above
+    --wide_file {WIDE_FILE_PATH} \
+    --lmdb_folder {LMDB_FOLDER_PATH} \
+    --splits_file {SPLITS_JSON_FILE_PATH} \
+    --selected_views A4C \ # see model_zoo/DROID/echo_defines.py for view/Doppler/quality/canonical axis arguments
+    --selected_views RV_focused \ 
+    --selected_doppler standard \
+    --selected_quality good \
+    --selected_canonical on_axis \
+    --n_train_patients all \
+    --batch_size 16\
+    --epochs 50 \
+    --es_patience 5 \
+    --scale_outputs \
+    --skip_modulo 4\
+    --adam 1e-4 \
+    --movinet_chkp_dir {MOVINET_CHECKPOINT_PATH} \ # corresponds to ml4h/model_zoo/DROID-RV/movinet_a2_base/chkp
+    --pretrained_chkp_dir {PRETRAINED_CHECKPOINT_PATH} \ # used when fine-tuning, for example can provide ml4h/model_zoo/DROID-RV/droid_rv_checkpoint/chkp
+    --output_dir {OUTPUT_FOLDER_PATH}
+```