PyTorch implementation of machine translation model in "Attention is All You Need".
Offical implementation in TensorFlow at here.
There are many other PyTorch implementations. This repo is different in:
-
Trained on a larger data set that is close to the original paper. (3M en-de sentence pairs, the original paper uses 4.5M).
-
Followed the original paper with byte pair encoding and label smoothing.
Dev data set: newstest2014+2015. Perplexity: 10.8.
Test data set: newstest2016. BLEU: 20.53.
from Model import Transformer
from Translator import Translator
# Load Transformer model and byte pair encoding model.
model = Transformer(bpe_size=32000, h=8, d_model=512, p=0.1, d_ff=1024).cuda()
model.load_state_dict(torch.load("epoch60_acc_58.90_ppl_10.85.pt")["model"])
translator = Translator(model, "en-de_bp3_32000.model")
# Translate one sentence.
target, attn = translator.translate("Every 10 minutes, bitcoin mining machines solve math problems that create 12.5 new bitcoin.")
print(target)
# > Jede zehn Minuten, Bitcoinbergbaumaschinen lösen Mathematikprobleme, die 12,5 neue Bitcoin erzeugen.
# Plot attention heatmap (average of 8 heads).
translator.attention_heatmap(attn)
See demo notebook.
-
Memory footprint is high. The largest batch size I used is only 96 with 11GB memory.
-
No support for multi-GPU yet.
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent