To train your model using mixed or TF32 precision with Tensor Cores or using FP32, perform the following steps using the default parameters of the GNMT v2 model on the WMT16 English German dataset. For the specifics concerning training and inference, see the Advanced section.

1. Clone the repository.

git clone https://github.com/NVIDIA/DeepLearningExamples
cd DeepLearningExamples/PyTorch/Translation/GNMT

2. Build the GNMT v2 Docker container.

bash scripts/docker/build.sh

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

bash scripts/docker/interactive.sh

4. Download and preprocess the dataset.

Data will be downloaded to the data directory (on the host). The data directory is mounted to the /workspace/gnmt/data location in the Docker container.

bash scripts/wmt16_en_de.sh

5. Start training.

The training script saves only one checkpoint with the lowest value of the loss function on the validation dataset. All results and logs are saved to the gnmt directory (on the host) or to the /workspace/gnmt/gnmt directory (in the container). By default, the train.py script will launch mixed precision training with Tensor Cores. You can change this behavior by setting:

• the --math fp32 flag to launch single precision training (for NVIDIA Volta and NVIDIA Turing architectures) or
• the --math tf32 flag to launch TF32 training with Tensor Cores (for NVIDIA Ampere architecture)

for the train.py training script.

To launch mixed precision training on 1, 4 or 8 GPUs, run:

python3 -m torch.distributed.launch --nproc_per_node=<#GPUs> train.py --seed 2 --train-global-batch-size 1024

To launch mixed precision training on 16 GPUs, run:

python3 -m torch.distributed.launch --nproc_per_node=16 train.py --seed 2 --train-global-batch-size 2048

By default, the training script will launch training with batch size 128 per GPU. If --train-global-batch-size is specified and larger than 128 times the number of GPUs available for the training then the training script will accumulate gradients over consecutive iterations and then perform the weight update. For example, 1 GPU training with --train-global-batch-size 1024 will accumulate gradients over 8 iterations before doing the weight update with accumulated gradients.

6. Start evaluation.

The training process automatically runs evaluation and outputs the BLEU score after each training epoch. Additionally, after the training is done, you can manually run inference on the test dataset with the checkpoint saved during the training.

To launch FP16 inference on the newstest2014.en test set, run:

python3 translate.py \
  --input data/wmt16_de_en/newstest2014.en \
  --reference data/wmt16_de_en/newstest2014.de \
  --output /tmp/output \
  --model gnmt/model_best.pth

The script will load the checkpoint specified by the --model option, then it will launch inference on the file specified by the --input option, and compute BLEU score against the reference translation specified by the --reference option. Outputs will be stored to the location specified by the --output option.

Additionally, one can pass the input text directly from the command-line:

python3 translate.py \
  --input-text "The quick brown fox jumps over the lazy dog" \
  --model gnmt/model_best.pth

Translated output will be printed to the console:

(...)
0: Translated output:
Der schnelle braune Fuchs springt über den faulen Hund

By default, the translate.py script will launch FP16 inference with Tensor Cores. You can change this behavior by setting:

• the --math fp32 flag to launch single precision inference (for NVIDIA Volta and NVIDIA Turing architectures) or
• the --math tf32 flag to launch TF32 inference with Tensor Cores (for NVIDIA Ampere architecture)

for the translate.py inference script.