To pretrain or fine tune your model for BioMedical tasks using mixed precision with Tensor Cores or using FP32, perform the following steps using the default parameters of the BERT model.

1. Clone the repository.

git clone https://github.com/NVIDIA/DeepLearningExamples
cd DeepLearningExamples/TensorFlow/LanguageModeling/BERT

2. Build the BERT TensorFlow NGC container.

bash scripts/docker/build.sh

3. Download and preprocess the PubMed dataset.

To download and preprocess pre-training data as well as the required vocab files, run the following script:

bash biobert/scripts/biobert_data_download.sh

Datasets for finetuning for NER can be obtained from this repository Datasets for finetuning for RE can be obtained from this repository

Place them both in /workspace/bert/data/biobert/ to be automatically picked up by our scripts.

4. Start an interactive session in the NGC container to run training/inference.

After you build the container image and download the data, you can start an interactive CLI session as follows:

bash scripts/docker/launch.sh

5. Download the pre-trained checkpoint, vocabulary, and configuration files.

We have uploaded checkpoints for fine tuning and pre-training on BioMedical Corpus's on the NGC Model Registry. You can download them directly from the NGC model catalog.

Place our BioBERT checkpoints in the results/ to easily access it in your scripts.

6. Start pre-training.

From within the container, you can use the following script to run the 1st phase of the pre-training using cased vocabulary:

bash biobert/scripts/run_pretraining-pubmed_base_phase_1.sh <train_batch_size> <learning_rate> <cased> <precision> <use_xla> <num_gpus> <warmup_steps> <train_steps> <num_accumulation_steps> <save_checkpoint_steps> <eval_batch_size>

For the 2nd phase of the pre-training, issue:

bash biobert/scripts/run_pretraining-pubmed_base_phase_2.sh <path_to_phase_1_checkpoint> <train_batch_size> <learning_rate> <cased> <precision> <use_xla> <num_gpus> <warmup_steps> <train_steps> <num_accumulation_steps> <save_checkpoint_steps> <eval_batch_size>

Refer to (MultiNode Section)[multi-node] for details on utilizing multiple nodes for faster pretraining.

6. Start fine tuning.

The above pretrained BERT representations can be fine tuned with just one additional output layer for a state-of-the-art biomedical text-mining system. From within the container, you can use the following script to run fine-training for NER.

Note: The scripts assume you are running on 16 V100 32GB GPUs. If you are running on GPU having less than 32GB memory or fewer GPUs, batch size, learning rate and number of GPUs needs to be adjusted.

For NER on disease entities:

bash biobert/scripts/ner_bc5cdr-disease.sh  <init_checkpoint> <train_batch_size> <learning_rate> <cased> <precision> <use_xla> <num_gpu> <seq_length> <bert_model> <eval_batch_size> <epochs>

For NER on chemical entities:

bash biobert/scripts/ner_bc5cdr-chem.sh  <init_checkpoint> <train_batch_size> <learning_rate> <cased> <precision> <use_xla> <num_gpu> <seq_length> <bert_model> <eval_batch_size> <epochs>

For relation extraction, issue:

bash biobert/scripts/rel_chemprot.sh <init_checkpoint> <train_batch_size> <learning_rate> <cased> <precision> <use_xla> <num_gpu> <seq_length> <bert_model> <eval_batch_size> <epochs>

8. Start validation/evaluation.

The biobert/scripts/run_biobert_finetuning_inference.sh script runs inference on a checkpoint fine tuned for a specific task and evaluates the validity of predictions on the basis of F1, precision and recall scores.

bash biobert/scripts/run_biobert_finetuning_inference.sh <task> <init_checkpoint> <bert_model> <cased> <precision> <use_xla> <batch_size>

For FP16 inference for NER on BC5DR Chemical task with XLA using a DGX-2 V100 32G, run:

bash biobert/scripts/run_biobert_finetuning_inference.sh ner_bc5cdr-chem /results/model.ckpt base false fp16 true 16

Tasks ner_bc5cdr-chem, ner_bc5cdr-disease and rel_chemprot are currently supported.