The following sections provide greater details of the dataset, running training and inference, and the training results.

Scripts and sample code

In the root directory, the most important files are:

• nmt.py: serves as the entry point to launch the training
• Dockerfile: container with the basic set of dependencies to run GNMT v2
• requirements.txt: set of extra requirements for running GNMT v2
• attention_wrapper.py, gnmt_model.py, model.py: model definition
• estimator.py: functions for training and inference

In the script directory, the most important files are:

• translate.py: wrapped on nmt.py for benchmarking and running inference
• parse_log.py: script for retrieving information in JSON format from the training log
• wmt16_en_de.sh: script for downloading and preprocessing the dataset

In the script/docker directory, the files are:

• build.sh: script for building the GNMT container
• interactive.sh: script for running the GNMT container interactively

Parameters

The most useful arguments are as follows:

  --learning_rate LEARNING_RATE
                        Learning rate.
  --warmup_steps WARMUP_STEPS
                        How many steps we inverse-decay learning.
  --max_train_epochs MAX_TRAIN_EPOCHS
                        Max number of epochs.
  --target_bleu TARGET_BLEU
                        Target bleu.
  --data_dir DATA_DIR   Training/eval data directory.
  --translate_file TRANSLATE_FILE
                        File to translate, works only with translate mode
  --output_dir OUTPUT_DIR
                        Store log/model files.
  --batch_size BATCH_SIZE
                        Total batch size.
  --log_step_count_steps LOG_STEP_COUNT_STEPS
                        The frequency, in number of global steps, that the
                        global step and the loss will be logged during training
  --num_gpus NUM_GPUS   Number of gpus in each worker.
  --random_seed RANDOM_SEED
                        Random seed (>0, set a specific seed).
  --ckpt CKPT           Checkpoint file to load a model for inference.
                        (defaults to newest checkpoint)
  --infer_batch_size INFER_BATCH_SIZE
                        Batch size for inference mode.
  --beam_width BEAM_WIDTH
                        beam width when using beam search decoder. If 0, use
                        standard decoder with greedy helper.
  --amp                 use amp for training and inference
  --mode {train_and_eval,infer,translate}

Command-line options

To see the full list of available options and their descriptions, use the -h or --help command line option, for example:

python nmt.py --help

Getting the data

The GNMT v2 model was trained on the WMT16 English-German dataset and newstest2014 is used as a testing dataset.

This repository contains the scripts/wmt16_en_de.sh download script which automatically downloads and preprocesses the training and test datasets. By default, data is downloaded to the data directory.

Our download script is very similar to the wmt16_en_de.sh script from the tensorflow/nmt repository. Our download script contains an extra preprocessing step, which discards all pairs of sentences which can't be decoded by latin-1 encoder. The scripts/wmt16_en_de.sh script uses the subword-nmt package to segment text into subword units (Byte Pair Encodings - BPE). By default, the script builds the shared vocabulary of 32,000 tokens.

In order to test with other datasets, the scripts need to be customized accordingly.

Dataset guidelines

The process of downloading and preprocessing the data can be found in the scripts/wmt16_en_de.sh script.

Initially, data is downloaded from www.statmt.org. Then, europarl-v7, commoncrawl and news-commentary corpora are concatenated to form the training dataset, similarly newstest2015 and newstest2016 are concatenated to form the validation dataset. Raw data is preprocessed with Moses, first by launching Moses tokenizer (tokenizer breaks up text into individual words), then by launching clean-corpus-n.perl which removes invalid sentences and does initial filtering by sequence length.

Second stage of preprocessing is done by launching the scripts/filter_dataset.py script, which discards all pairs of sentences that can't be decoded by latin-1 encoder.

Third state of preprocessing uses the subword-nmt package. First it builds shared byte pair encoding vocabulary with 32,000 merge operations (command subword-nmt learn-bpe), then it applies generated vocabulary to training, validation and test corpora (command subword-nmt apply-bpe).

Training process

The training configuration can be launched by running the nmt.py script. By default, the training script saves the checkpoint after every training epoch and after every 2000 training steps within each epoch. Results are stored in the results directory.

The training script launches data-parallel training on multiple GPUs. We have tested reliance on up to 8 GPUs on a single node.

After each training epoch, the script runs an evaluation and outputs a BLEU score on the test dataset (newstest2014). BLEU is computed by the SacreBLEU package. Logs from the training and evaluation are saved to the results directory.

The training script automatically runs testing after each training epoch. The results from the testing are printed to the standard output and saved to the log files.

The summary after each training epoch is printed in the following format:

training time for epoch 1: 29.37 mins (2918.36 sent/sec, 139640.48 tokens/sec)
[...]
bleu is 20.50000
eval time for epoch 1: 1.57 mins (78.48 sent/sec, 4283.88 tokens/sec)

The BLEU score is computed on the test dataset. Performance is reported in total sentences per second and in total tokens per second. The performance result is averaged over an entire training epoch and summed over all GPUs participating in the training.

To view all available options for training, run python nmt.py --help.

Inference process

Validation and translation can be run by launching the nmt.py script, although, it requires a pre-trained model checkpoint and tokenized input (for validation) and non-tokenized input (for translation).

Validation process

The nmt.py script supports batched validation (--mode=infer flag). By default, it launches beam search with beam size of 5, coverage penalty term and length normalization term. Greedy decoding can be enabled by setting the --beam_width=1 flag for the nmt.py inference script. To control the batch size use the --infer_batch_size flag.

To view all available options for validation, run python nmt.py --help.

Translation process

The nmt.py script supports batched translation (--mode=translate flag). By default, it launches beam search with beam size of 5, coverage penalty term and length normalization term. Greedy decoding can be enabled by setting the --beam_width=1 flag for the nmt.py prediction script. To control the batch size use the --infer_batch_size flag.

The input file may contain many sentences, each on a new line. The file can be specified by the --translate_file <file> flag. This script will create a new file called <file>.trans, with translation of the input file.

To view all available options for translation, run python nmt.py --help.