This implementation of Transformer model architecture is based on the optimized implementation in Fairseq NLP toolkit.
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 Transformer model on the WMT14 English-German dataset. For the specifics concerning training and inference, see the Advanced section.
- Clone the repository
git clone https://github.com/NVIDIA/DeepLearningExamples.git
cd DeepLearningExamples/PyTorch/Translation/Transformer
- Build and launch the Transformer PyTorch NGC container
docker build . -t your.repository:transformer
nvidia-docker run -it --rm --ipc=host your.repository:transformer bash
If you already have preprocessed data, use:
nvidia-docker run -it --rm --ipc=host -v <path to your preprocessed data>:/data/wmt14_en_de_joined_dict your.repository:transformer bash
If you already have data downloaded, but it has not yet been preprocessed, use:
nvidia-docker run -it --rm --ipc=host -v <path to your unprocessed data>:/workspace/translation/examples/translation/orig your.repository:transformer bash
- Download and preprocess dataset: Download and preprocess the WMT14 English-German dataset.
scripts/run_preprocessing.sh
After running this command, data will be downloaded to /workspace/translation/examples/translation/orig directory and this data will be processed and put into /data/wmt14_en_de_joined_dict directory.
- Start training
python -m torch.distributed.launch --nproc_per_node 8 /workspace/translation/train.py /data/wmt14_en_de_joined_dict \
--arch transformer_wmt_en_de_big_t2t \
--share-all-embeddings \
--optimizer adam \
--adam-betas '(0.9, 0.997)' \
--adam-eps "1e-9" \
--clip-norm 0.0 \
--lr-scheduler inverse_sqrt \
--warmup-init-lr 0.0 \
--warmup-updates 4000 \
--lr 0.0006 \
--min-lr 0.0 \
--dropout 0.1 \
--weight-decay 0.0 \
--criterion label_smoothed_cross_entropy \
--label-smoothing 0.1 \
--max-tokens 5120 \
--seed 1 \
--fuse-layer-norm \
--amp \
--amp-level O2 \
--save-dir /workspace/checkpoints \
--distributed-init-method env://
The script saves checkpoints every epoch to the directory specified in the --save-dir option. In addition, the best performing checkpoint (in terms of loss) and the latest checkpoints are saved separately.
WARNING: If you don't have access to sufficient disk space, use the --save-interval $N option. The checkpoints are ~3.4GB large. For example, it takes the Transformer model 30 epochs for the validation loss to plateau. The default option is to save last checkpoint, the best checkpoint and a checkpoint for every epoch, which means (30+1+1)*3.4GB = 108.8GB of a disk space used. Specifying --save-interval 10 reduces this to (30/10+1+1)*3.4GB = 17GB.
- Start interactive inference
python inference.py \
--buffer-size 5000 \
--path /path/to/your/checkpoint.pt \
--max-tokens 10240 \
--fuse-dropout-add \
--remove-bpe \
--bpe-codes /path/to/bpe_code_file \
--fp16
where,
--pathoption is the location of the checkpoint file.--bpe-codesoption is the location of thecodefile. If the default training command mentioned above is used, this file can be found in the preprocessed data ( i.e.,/data/wmt14_en_de_joined_dict) directory.