NVIDIA

Parabricks DeepVariant Retraining Notebook

Resource

NVIDIA

Parabricks DeepVariant Retraining Notebook

This notebook shows how to retrain DeepVariant models using Parabricks.

Back to File Browser

Retraining_DeepVariant.ipynb

Retraining DeepVariant using Parabricks

In this lab we will use Parabricks and DeepVariant to retrain a DeepVariant model on new data.

What is Parabricks?

NVIDIA Parabricks® is the only GPU-accelerated computational genomics toolkit that delivers fast and accurate analysis for sequencing centers, clinical teams, genomics researchers, and next-generation sequencing instrument developers. Parabricks provides GPU-accelerated versions of tools used every day by computational biologists and bioinformaticians—enabling significantly faster runtimes, workflow scalability, and lower compute costs.

The toolkit includes full compatibility with workflow languages and managers (WDL, NextFlow, Cromwell) to easily intertwine GPU- and CPU-powered tasks, as well as support for easy cloud deployment (AWS, GCP, Terra, and DNAnexus).

What is DeepVariant?

DeepVariant, developed by Google, is a deep learning-based variant caller that takes aligned reads, produces pileup image tensors from them, classifies each tensor using a convolutional neural network, and then outputs the results in a VCF or gVCF file.

Package dependencies

There are a few tools we need to install to successfully run this lab. They can be installed by running the cells below.

To run a cell, you can either click into it and hit Ctrl+Enter or you can hit the Run button on the top toolbar.

In [ ]:

%%sh

# install google packages to access and download data
sudo apt-get install -y apt-transport-https ca-certificates gnupg
echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | sudo tee -a /etc/apt/sources.list.d/google-cloud-sdk.list
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key --keyring /usr/share/keyrings/cloud.google.gpg add -
sudo apt-get update && sudo apt install -y google-cloud-cli

In [ ]:

%%sh

# download extra packages
sudo apt install -y parallel
sudo apt install -y python3-pip && python3 -m pip install tensorflow==2.11.0

DeepVariant baseline performance

DeepVariant comes with a generic, baseline model for WGS data. Let's see the baseline performance on our test data (Chromosome 20 from HG001).

In [ ]:

%%sh 

# Let's make these folder paths (input/data should already exist)
mkdir -p output/data
mkdir -p output/logs

In [ ]:

%%sh 

BIN_VERSION="1.4.0"
INPUT_DIR="input/data"
OUTPUT_DIR="output/data"
REF="${INPUT_DIR}/ucsc_hg19.fa"
BAM_CHR20="${INPUT_DIR}/BGISEQ_PE100_NA12878.sorted.chr20.bam"

docker run --gpus all \
    -v `pwd`:`pwd` \
    -w `pwd` \
    google/deepvariant:"${BIN_VERSION}-gpu" \
    /opt/deepvariant/bin/run_deepvariant \
    --model_type WGS \
    --ref "${REF}" \
    --reads "${BAM_CHR20}" \
    --regions "chr20" \
    --output_vcf "${OUTPUT_DIR}/baseline.vcf.gz" \
    --num_shards=16

Let's compare the baseline VCF to the groundtruth VCF and see what the accuracy looks like.

In [ ]:

%%sh 

# RUNTIME ~ 3 minutes 

INPUT_DIR="input/data"
OUTPUT_DIR="output/data"

REF="${INPUT_DIR}/ucsc_hg19.fa"
TRUTH_VCF="${INPUT_DIR}/HG001_GRCh37_GIAB_highconf_CG-IllFB-IllGATKHC-Ion-10X-SOLID_CHROM1-X_v.3.3.2_highconf_PGandRTGphasetransfer_chrs_FIXED.vcf.gz"
TRUTH_BED="${INPUT_DIR}/HG001_GRCh37_GIAB_highconf_CG-IllFB-IllGATKHC-Ion-10X-SOLID_CHROM1-X_v.3.3.2_highconf_nosomaticdel_chr.bed"

time docker run \
    -v `pwd`:`pwd` \
    -w `pwd` \
    jmcdani20/hap.py:v0.3.12 /opt/hap.py/bin/hap.py \
    "${TRUTH_VCF}" \
    "${OUTPUT_DIR}/baseline.vcf.gz" \
    -f "${TRUTH_BED}" \
    -r "${REF}" \
    -o "${OUTPUT_DIR}/chr20-calling.happy.output" \
    -l chr20 \
    --engine=vcfeval \
    --pass-only > "baseline_hap.log" 2>&1

Let's look at the accuracy of the baseline WGS model on our dataset. For this, we can cat the baseline_hap.log file to see what hap.py generated.

In [ ]:

! cat baseline_hap.log

The performance here is pretty good. The goal of the baseline WGS model is to do well across a wide spectrum of data. However, we can retrain this WGS model to perform better for our specific use case.

Building a dataset to retrain DeepVariant

To retrain the WGS baseline model, we need a dataset to train on.

We will use Chromosome 1 from HG001 to generate a training dataset.

In [ ]:

%%sh

# RUNTIME ~ 5 mins

INPUT_DIR="input/data"
OUTPUT_DIR="output/data"
REF="${INPUT_DIR}/ucsc_hg19.fa"
BAM_CHR1="${INPUT_DIR}/BGISEQ_PE100_NA12878.sorted.chr1.bam"
TRUTH_VCF="${INPUT_DIR}/HG001_GRCh37_GIAB_highconf_CG-IllFB-IllGATKHC-Ion-10X-SOLID_CHROM1-X_v.3.3.2_highconf_PGandRTGphasetransfer_chrs_FIXED.vcf.gz"
TRUTH_BED="${INPUT_DIR}/HG001_GRCh37_GIAB_highconf_CG-IllFB-IllGATKHC-Ion-10X-SOLID_CHROM1-X_v.3.3.2_highconf_nosomaticdel_chr.bed"
TRAIN_EXAMPLES="${OUTPUT_DIR}/training_set_gpu.with_label.tfrecord.gz"

docker run \
    --gpus all --rm \
    -v `pwd`:`pwd` \
    -w `pwd` \
    nvcr.io/nvidia/clara/deepvariant_train:4.1.0-1 \
    /usr/local/parabricks/binaries/bin/deepvariant \
    "${REF}" \
    "${BAM_CHR1}" \
    2 \
    -o "${OUTPUT_DIR}/output.vcf" \
    -n 8 \
    --channel_insert_size \
    -L "chr1" \
    -disable-use-window-selector-model \
    --mode training \
    --truth_variants "${TRUTH_VCF}" \
    --confident_regions "${TRUTH_BED}" \
    --examples "${TRAIN_EXAMPLES}" \
    -z 4

We will use Chromosome 21 from HG001 to generate a validation dataset.

In [ ]:

%%sh 

# RUNTIME ~ 2 mins 

INPUT_DIR="input/data"
OUTPUT_DIR="output/data"
REF="${INPUT_DIR}/ucsc_hg19.fa"
BAM_CHR21="${INPUT_DIR}/BGISEQ_PE100_NA12878.sorted.chr21.bam"
TRUTH_VCF="${INPUT_DIR}/HG001_GRCh37_GIAB_highconf_CG-IllFB-IllGATKHC-Ion-10X-SOLID_CHROM1-X_v.3.3.2_highconf_PGandRTGphasetransfer_chrs_FIXED.vcf.gz"
TRUTH_BED="${INPUT_DIR}/HG001_GRCh37_GIAB_highconf_CG-IllFB-IllGATKHC-Ion-10X-SOLID_CHROM1-X_v.3.3.2_highconf_nosomaticdel_chr.bed"
VAL_EXAMPLES="${OUTPUT_DIR}/validation_set_gpu.with_label.tfrecord.gz"

docker run \
    --gpus all --rm \
    -v `pwd`:`pwd` \
    -w `pwd` \
    nvcr.io/nvidia/clara/deepvariant_train:4.1.0-1 \
    /usr/local/parabricks/binaries/bin/deepvariant \
    "${REF}" \
    "${BAM_CHR21}" \
    2 \
    -o "${OUTPUT_DIR}/output.vcf" \
    -n 8 \
    --channel_insert_size \
    -L "chr21" \
    -disable-use-window-selector-model \
    --mode training \
    --truth_variants "${TRUTH_VCF}" \
    --confident_regions "${TRUTH_BED}" \
    --examples "${EXAMPLES}" \
    -z 4

Shuffling the training data

Before we can train the model we will need to shuffle each set of examples and generate a data config file. This has to be done for both the training and validation dataset.

In [ ]:

%%sh

OUTPUT_DIR="output/data"

# RUNTIME ~ 25 minutes 

# shuffle training set 
time python3 scripts/shuffle_tfrecords_lowmem.py \
    --input_pattern_list="${OUTPUT_DIR}/training_set_gpu.with_label.tfrecord-?????-of-00004.gz" \
    --output_pattern_prefix="${OUTPUT_DIR}/training_set_gpu.with_label.shuffled" \
    --output_dataset_config="${OUTPUT_DIR}/training_set_gpu.pbtxt" \
    --output_dataset_name="HG001" \
    --direct_num_workers=8 \
    --step=-1

In [ ]:

%%sh

# RUNTIME ~ 4 minutes 

OUTPUT_DIR="output/data"

# shuffle validation set
time python3 scripts/shuffle_tfrecords_lowmem.py \
    --input_pattern_list="${OUTPUT_DIR}/validation_set_gpu.with_label.tfrecord-?????-of-00004.gz" \
    --output_pattern_prefix="${OUTPUT_DIR}/validation_set_gpu.with_label.shuffled" \
    --output_dataset_config="${OUTPUT_DIR}/validation_set_gpu.pbtxt" \
    --output_dataset_name="HG001" \
    --direct_num_workers=8 \
    --step=-1

Training the DeepVariant model

Next we want to run the following two code blocks at the same time to train and evaluate the different possible models.

This first cell will constantly check the training_dir folder for new model checkpoints. When a new model checkpoint is generated by the training script, it will evaluate the checkpoint and keep track of which checkpoint performs the best.

In [ ]:

%%sh

# have to manually stop running this once model train stops generating checkpoints

BIN_VERSION="1.4.0"
OUTPUT_DIR="output/data"
TRAINING_DIR="training_dir"
LOG_DIR="output/logs"

docker run \
    -v `pwd`:`pwd` \
    -w `pwd` \
    google/deepvariant:"${BIN_VERSION}" \
    /opt/deepvariant/bin/model_eval \
    --dataset_config_pbtxt="${OUTPUT_DIR}/validation_set_gpu.pbtxt" \
    --checkpoint_dir="${TRAINING_DIR}" \
    --batch_size=512 > "${LOG_DIR}/eval_gpu.log" 2>&1 &

This cell kicks off the DeepVariant training.

In [ ]:

%%sh

# all parameters below are used as an example. They are not optimized for this dataset, and are not recommended as the best default

BIN_VERSION="1.4.0"
OUTPUT_DIR="output/data"
LOG_DIR="output/logs"
TRAINING_DIR="training_dir"

MODEL_BUCKET="gs://deepvariant/models/DeepVariant/${BIN_VERSION}/DeepVariant-inception_v3-${BIN_VERSION}+data-wgs_standard"
GCS_PRETRAINED_WGS_MODEL="${MODEL_BUCKET}/model.ckpt"

(docker run \
    --gpus all \
    -v `pwd`:`pwd` \
    -w `pwd` \
    google/deepvariant:"${BIN_VERSION}-gpu" \
    /opt/deepvariant/bin/model_train \
    --dataset_config_pbtxt="${OUTPUT_DIR}/training_set_gpu.pbtxt" \
    --train_dir="${TRAINING_DIR}" \
    --model_name="inception_v3" \
    --number_of_steps=5000 \
    --save_interval_secs=300 \
    --batch_size=32 \
    --learning_rate=0.0005 \
    --start_from_checkpoint="${GCS_PRETRAINED_WGS_MODEL}" \
    ) > "${LOG_DIR}/train_gpu.log" 2>&1

Choose the best model

We then want to pick the best mdoel. We can determine which model to use by running the line of code below.

Let's cat the file training_dir/best_checkpoint.txt

In [ ]:

! cat training_dir/best_checkpoint.txt

Re-run DeepVariant using this new model

Now that we know which model checkpoint performed the best, we can use this model to run DeepVariant on Chromosome 20 again, but this time, swap out the BEST_CHECKPOINT variable for your best performing model checkpoint.

In [ ]:

%%sh

# RUNTIME ~ 10 minutes 

BIN_VERSION="1.4.0"
TRAINING_DIR="training_dir"
INPUT_DIR="input/data"
OUTPUT_DIR="output/data"
BEST_CHECKPOINT="model.ckpt-0" # CHANGE THIS LINE 

REF="${INPUT_DIR}/ucsc_hg19.fa"
BAM_CHR20="${INPUT_DIR}/BGISEQ_PE100_NA12878.sorted.chr20.bam"

# model to train
time docker run \
    --gpus all \
    -v `pwd`:`pwd` \
    -w `pwd` \
    google/deepvariant:"${BIN_VERSION}-gpu" \
    /opt/deepvariant/bin/run_deepvariant \
    --model_type WGS \
    --customized_model "${TRAINING_DIR}/${BEST_CHECKPOINT}" \
    --ref "${REF}" \
    --reads "${BAM_CHR20}" \
    --regions "chr20" \
    --output_vcf "${OUTPUT_DIR}/model_set_gpu.vcf.gz" \
    --num_shards=16

Evaluate the performance

Lastly, we can evaluate the performance of our DeepVariant model using hap.py and compare to the performance to the baseline model.

In [ ]:

%%sh 

# RUNTIME ~ 3 minutes 

INPUT_DIR="input/data"
OUTPUT_DIR="output/data"

REF="${INPUT_DIR}/ucsc_hg19.fa"
TRUTH_VCF="${INPUT_DIR}/HG001_GRCh37_GIAB_highconf_CG-IllFB-IllGATKHC-Ion-10X-SOLID_CHROM1-X_v.3.3.2_highconf_PGandRTGphasetransfer_chrs_FIXED.vcf.gz"
TRUTH_BED="${INPUT_DIR}/HG001_GRCh37_GIAB_highconf_CG-IllFB-IllGATKHC-Ion-10X-SOLID_CHROM1-X_v.3.3.2_highconf_nosomaticdel_chr.bed"

time docker run \
    -v `pwd`:`pwd` \
    -w `pwd` \
    jmcdani20/hap.py:v0.3.12 /opt/hap.py/bin/hap.py \
    "${TRUTH_VCF}" \
    "${OUTPUT_DIR}/model_set_gpu.vcf.gz" \
    -f "${TRUTH_BED}" \
    -r "${REF}" \
    -o "${OUTPUT_DIR}/chr20-calling.happy.output" \
    -l chr20 \
    --engine=vcfeval \
    --pass-only > "retrained_hap.log" 2>&1

Let's cat the retrained_hap.log to see the accuracy metrics.

In [ ]:

! cat retrained_hap.log

Conclusion

Congratulations! You have now successfully retrained DeepVariant on new data and evaluated the performance compared to a baseline. Now you can take this knowledge and try it on your own data to see if you can get better results.