Counting 3,834 Big Data & Machine Learning Frameworks, Toolsets, and Examples...
Suggestion? Feedback? Tweet @stkim1

Attention is all you need: A Pytorch Implementation

This is a PyTorch implementation of the Transformer model in "Attention is All You Need" (Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, Illia Polosukhin, arxiv, 2017).

A novel sequence to sequence framework utilizes the self-attention mechanism, instead of Convolution operation or Recurrent structure, and achieve the state-of-the-art performance on WMT 2014 English-to-German translation task. (2017/06/12)

The official Tensorflow Implementation can be found in: tensorflow/tensor2tensor.

To learn more about self-attention mechanism, you could read "A Structured Self-attentive Sentence Embedding".

The project support training and translation with trained model now.

Note that this project is still a work in progress.

If there is any suggestion or error, feel free to fire an issue to let me know. :)


  • python 3.4+
  • pytorch 0.4.1+
  • tqdm
  • numpy


Some useful tools:

The example below uses the Moses tokenizer ( to prepare the data and the moses BLEU script for evaluation.

sed -i "s/$RealBin\/..\/share\/nonbreaking_prefixes//" tokenizer.perl

WMT'16 Multimodal Translation: Multi30k (de-en)

An example of training for the WMT'16 Multimodal Translation task (

0) Download the data.

mkdir -p data/multi30k
wget &&  tar -xf training.tar.gz -C data/multi30k && rm training.tar.gz
wget && tar -xf validation.tar.gz -C data/multi30k && rm validation.tar.gz
wget && tar -xf mmt16_task1_test.tar.gz -C data/multi30k && rm mmt16_task1_test.tar.gz

1) Preprocess the data.

for l in en de; do for f in data/multi30k/*.$l; do if [[ "$f" != *"test"* ]]; then sed -i "$ d" $f; fi;  done; done
for l in en de; do for f in data/multi30k/*.$l; do perl tokenizer.perl -a -no-escape -l $l -q  < $f > $f.atok; done; done
python -train_src data/multi30k/train.en.atok -train_tgt data/multi30k/ -valid_src data/multi30k/val.en.atok -valid_tgt data/multi30k/ -save_data data/

2) Train the model

python -data data/ -save_model trained -save_mode best -proj_share_weight -label_smoothing

If your source and target language share one common vocabulary, use the -embs_share_weight flag to enable the model to share source/target word embedding.

3) Test the model

python -model trained.chkpt -vocab data/ -src data/multi30k/test.en.atok -no_cuda



  • Parameter settings:

    • default parameter and optimizer settings
    • label smoothing
    • target embedding / pre-softmax linear layer weight sharing.
  • Elapse per epoch (on NVIDIA Titan X):

    • Training set: 0.888 minutes
    • Validation set: 0.011 minutes


  • coming soon.


  • Evaluation on the generated text.
  • Attention weight plot.


  • The project structure, some scripts and the dataset preprocessing steps are heavily borrowed from OpenNMT/OpenNMT-py.
  • Thanks for the suggestions from @srush, @iamalbert and @ZiJianZhao.