3 - Neural Machine Translation by Jointly Learning to Align and Translate

我们将根据论文[Neural Machine Translation by Jointly Learning to Align and Translate]实现模型（https://arxiv.org/abs/1409.0473）。这个模型达到了最好的困惑度，约为27，而之前的模型约为34。

Introduction

在之前的模型中，我们为了减少信息压缩，在每一步中都显示地将上下文向量(context vector)和目标token(target sequence)传递给解码器，最后把目标输入$ y_t $以及隐藏状态$ s_t $和上下文向量$ z $到线性层$ f $进行预测。模型的示意图如下：

decoder

即使我们减少了一些压缩，我们的上下文向量仍然需要包含有关源句子的所有信息。在本篇博文中实现的模型将通过允许解码器在解码的每一个步骤都可以查看整个源句（通过其隐藏状态）来避免这种压缩！它是如何做到的？它使用注意力机制(attention)。

注意力机制首先计算的是注意力向量$ a $，它有着源语句的长度。 注意力向量每个元素的取值在0~1，整个向量的所有元素总和为1。然后我们计算源语句隐藏状态的加权和，$ H $，以获得加权源向量(weighted source vector)，$ w $。

$w = \sum_{i}a_ih_i$

我们在解码时的每个时间步骤都计算新的加权源向量，使用它作为我们的解码器RNN的输入以及线性层的输入来进行预测。我们将在接下来的内容中介绍。

Preparing Data

同样，准备工作与之前的模型类似。

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

from torchtext.datasets import TranslationDataset, Multi30k
from torchtext.data import Field, BucketIterator

import spacy

import random
import math
import os
import time
SEED = 1

random.seed(SEED)
torch.manual_seed(SEED)
torch.backends.cudnn.deterministic = True

spacy_de = spacy.load('de')
spacy_en = spacy.load('en')

def tokenize_de(text):
    """
    Tokenizes German text from a string into a list of strings
    """
    return [tok.text for tok in spacy_de.tokenizer(text)]

def tokenize_en(text):
    """
    Tokenizes English text from a string into a list of strings
    """
    return [tok.text for tok in spacy_en.tokenizer(text)]
    
SRC = Field(tokenize=tokenize_de, init_token='<sos>', eos_token='<eos>', lower=True)
TRG = Field(tokenize=tokenize_en, init_token='<sos>', eos_token='<eos>', lower=True)

train_data, valid_data, test_data = Multi30k.splits(exts=('.de', '.en'), fields=(SRC, TRG))

SRC.build_vocab(train_data, min_freq=2)
TRG.build_vocab(train_data, min_freq=2)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
BATCH_SIZE = 128

train_iterator, valid_iterator, test_iterator = BucketIterator.splits(
    (train_data, valid_data, test_data), batch_size=BATCH_SIZE, device=device)

Building the Seq2Seq Model

Encoder

首先，我们将构建编码器。与之前的模型类似，我们只使用单层GRU，但我们现在使用双向RNN。使用双向RNN，每一层将会有两个RNN模块。一个前向RNN(forward) 从左到右遍历句子（如下绿色所示），一个后向RNN从右到左遍历句子（黄色）。我们在代码中需要做的就是设置bidirectional = True，然后像以前一样将嵌入的句子传递给RNN。

encoder

现在我们有：

$\begin{align*} h_t^\rightarrow &= \text{EncoderGRU}^\rightarrow(x_t^\rightarrow,h_t^\rightarrow)\\ h_t^\leftarrow &= \text{EncoderGRU}^\leftarrow(x_t^\leftarrow,h_t^\leftarrow) \end{align*}$

其中$ x_0^\rightarrow = \text { <sos> }，x_1^\rightarrow = \text{guten} $和$ x_0^\leftarrow = \text {<eos>}，x_1^\leftarrow = \text {morgen}$。

和以前一样，我们只需将输入（embedded）传递给RNN，它告诉PyTorch将前向和后向初始隐藏状态（分别为$ h_0 ^\rightarrow $和$ h_0^\leftarrow $）初始化为所有元素都为零的张量。我们还将获得两个上下文向量，一个来自前向RNN，它看到句子中的最后一个单词后，$ z^\rightarrow = h_T^\rightarrow $；一个来自后向RNN，看到句子中第一个单词，$ z^\leftarrow = h_T^\leftarrow $。

最后，RNN返回outputs和hidden。

outputs的大小为[src sent len, batch size, hid dim * num directions]其中第三个轴中的第一个hid_dim元素是来自顶层中前向RNN的隐藏状态，而最后一个 hid_dim元素是顶层后向RNN的隐藏状态。你可以将第三轴视为堆叠在彼此顶部的前向和后向隐藏状态，即$ h_1 = [h_1^\rightarrow; h_ {T}^\leftarrow] $，$ h_2 = [h_2^\rightarrow; h_ {T-1}^\leftarrow] $。我们可以将所有堆叠的编码器隐藏状态表示为$ H = {h_1，h_2，…，h_T } $。

hidden的大小为[n layers * num direction，batch size，hid dim]，其中[ -2，:，:]表示最后一步的顶层前向RNN隐藏状态（即在看到句子中的最后一个单词之后）,[ - 1，：，：]表示最后一步的顶层后向RNN隐藏状态 (即在看到第一个单词之后)。

由于解码器不是双向的，它只需要一个上下文向量$ z $作为其初始隐藏状态$ s_0 $。我们目前有两个，一个前向和一个后向（$ z^\rightarrow = h_T^\rightarrow $和$ z^\leftarrow = h_T^\leftarrow $）。我们通过将两个上下文向量连接在一起，然后通过线性层$ g $并应用$ \ tanh $激活函数来解决这个问题。

$z=\tanh(g(h_T^\rightarrow, h_T^\leftarrow)) = \tanh(g(z^\rightarrow, z^\leftarrow)) = s_0$

注意：这实际上是与论文里的方法不同。相反，论文里仅仅将后向RNN隐藏状态通过线性层变换以获得上下文向量/解码器初始隐藏状态。这对我来说似乎没有意义，所以我改变了它。

因为希望我们的模型回顾整个源句，我们返回outputs——源语句中每个token堆叠的前向和后向隐藏状态。我们还返回hidden，它在解码器中充当我们的初始隐藏状态。

class Encoder(nn.Module):
    def __init__(self, input_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout):
        super().__init__()
        
        self.input_dim = input_dim
        self.emb_dim = emb_dim
        self.enc_hid_dim = enc_hid_dim
        self.dec_hid_dim = dec_hid_dim
        self.dropout = dropout
        
        self.embedding = nn.Embedding(input_dim, emb_dim)
        
        self.rnn = nn.GRU(emb_dim, enc_hid_dim, bidirectional=True)
        
        self.fc = nn.Linear(enc_hid_dim * 2, dec_hid_dim)
        
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, src):
        
        #src = [src sent len, batch size]
        
        embedded = self.dropout(self.embedding(src))
        
        #embedded = [src sent len, batch size, emb dim]
        
        outputs, hidden = self.rnn(embedded)
                
        #outputs = [src sent len, batch size, hid dim * num directions]
        #hidden = [n layers * num directions, batch size, hid dim]
        
        #hidden is stacked [forward_1, backward_1, forward_2, backward_2, ...]
        #outputs are always from the top layer
        
        #hidden [-2, :, : ] is the last of the forwards RNN 
        #hidden [-1, :, : ] is the last of the backwards RNN
        
        #initial decoder hidden is final hidden state of the forwards and backwards encoder RNNs fed through a linear layer
        hidden = torch.tanh(self.fc(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim=2)))
        
        #outputs = [src sent len, batch size, enc hid dim * 2]
        #hidden = [batch size, dec hid dim]
        
        return outputs, hidden

Attention

接下来是注意力层(attention layer)。这将采用解码器的上一步的隐藏状态$ s_ {t-1} $，以及来自编码器的所有堆叠的前向和后向隐藏状态，$ H $。该层将输出一个注意向量$ a_t $，它有着源句子的长度，每个元素在0和1之间，向量的所有元素总和为1。

直观地说，注意力层采用我们迄今已解码的$ s_ {t-1} $，以及我们编码的所有内容$ H $来生成一个向量$ a_t $，它表示为了正确预测要解码的下一个单词($ \hat{y}_{t + 1} $)，我们应该特别注意源语句的哪些单词。

首先，我们计算解码器的上一个隐藏状态和编码器隐藏状态之间的energy。由于我们的编码器隐藏状态是$ T $张量序列，而我们之前的解码器隐藏状态是单个张量，我们做的第一件事就是重复解码器的上一个隐藏状态$ T $次。然后我们通过将它们连接在一起并通过线性层（attn）和$ \tanh $激活函数来计算它们之间的能量$ E_t $。

$E_t = \tanh(\text{attn}(s_{t-1}, H))$

这可以被认为是计算编码器的每个隐藏状态和解码器的上一个隐藏状态“匹配”的程度。

我们目前batch中的每个例子都有[dec hid dim，src sent len]张量。我们希望这对于batch中的每个示例都是[src sent len]，因为注意力应该超过源句子的长度。这是通过将energy乘以[1，dec hid dim]张量，$ v $来实现的。

$\hat{a}_t = v E_t$

我们可以将此视为计算编码器的每个隐藏状态的所有dec_hid_dem元素的“匹配”加权和，其中权重是需要学习的参数（当我们学习$ v $的参数时）。

最后，我们确保注意力向量符合所有元素在0和1之间的约束，并且通过将它传递到$ \text{softmax} $层来保证所有元素的和为1。

$a_t = \text{softmax}(\hat{a_t})$

从图形上看，这看起来如下所示。这是用于计算第一个注意向量，其中$ s_ {t-1} = s_0 = z $。绿色/黄色块表示来自前向和后向RNN的隐藏状态，并且注意力计算全部在粉红色块内完成。

attention

class Attention(nn.Module):
    def __init__(self, enc_hid_dim, dec_hid_dim):
        super().__init__()
        
        self.enc_hid_dim = enc_hid_dim
        self.dec_hid_dim = dec_hid_dim
        
        self.attn = nn.Linear((enc_hid_dim * 2) + dec_hid_dim, dec_hid_dim)
        self.v = nn.Parameter(torch.rand(dec_hid_dim))
        
    def forward(self, hidden, encoder_outputs):
        
        #hidden = [batch size, dec hid dim]
        #encoder_outputs = [src sent len, batch size, enc hid dim * 2]
        
        batch_size = encoder_outputs.shape[1]
        src_len = encoder_outputs.shape[0]
        
        #repeat encoder/decoder hidden state src_len times
        hidden = hidden.unsqueeze(1).repeat(1, src_len, 1)
        
        encoder_outputs = encoder_outputs.permute(1, 0, 2)
        
        #hidden = [batch size, src sent len, dec hid dim]
        #encoder_outputs = [batch size, src sent len, enc hid dim * 2]
        
        energy = torch.tanh(self.attn(torch.cat((hidden, encoder_outputs), dim=2))) 
        
        #energy = [batch size, src sent len, dec hid dim]
        
        energy = energy.permute(0, 2, 1)
        
        #energy = [batch size, dec hid dim, src sent len]
        
        #v = [dec hid dim]
        
        v = self.v.repeat(batch_size, 1).unsqueeze(1)
        
        #v = [batch size, 1, dec hid dim]
                
        attention = torch.bmm(v, energy).squeeze(1)
        
        #attention= [batch size, src len]
        
        return F.softmax(attention, dim=1)

Decoder

接下来是解码器。

解码器包含注意力层 ——它读取解码器上一步的隐藏状态，$ s_{t-1} $和编码器的全部隐藏状态，$ H $，最后返回注意力向量$ a_t $。

然后我们使用这个注意力向量来创建一个加权源向量—— $ w_t $，用weighted表示，它是编码器隐藏状态的加权和，$ H $，使用$ a_t $作为权重。

$w_t = a_t H$

把输入(target source embedded) $ y_t $，加权源向量 $ w_t $和解码器的上一步隐藏状态$ s_ {t-1} $全部传递到解码器RNN中。在这里我们把$ y_t $和$ w_t $两个向量连在一起作为输入。

$s_t = \text{DecoderGRU}(y_t, w_t, s_{t-1})$

然后我们传递$ y_t $，$ w_t $和$ s_t $到线性层$ f $，去预测目标句子中的下一个单词$ \hat{y}_{t+1} $。这里我们将三个向量连接起来作为线性层$ f $的输入。

$\hat{y}_{t+1} = f(y_t, w_t, s_t)$

下图展示了如何解码示例中的第一个单词。

decoder

绿色/黄色代表$ H $的前向/后向编码器RNN，红色块代表上下文向量(context vector)——$ z = h_T = \tanh(g(h^\rightarrow_T，h^\leftarrow_T))= \tanh (g(z^\rightarrow，z^\leftarrow))= s_0 $，蓝色块代表输出$ s_t $的解码器RNN，紫色块代表线性层$ f $，输出$ \hat{y}_{t+1} $和橙色块代表加权源向量——$ w_t $。未显示的是注意力向量$ a_t $的计算。

class Decoder(nn.Module):
    def __init__(self, output_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout, attention):
        super().__init__()

        self.emb_dim = emb_dim
        self.enc_hid_dim = enc_hid_dim
        self.dec_hid_dim = dec_hid_dim
        self.output_dim = output_dim
        self.dropout = dropout
        self.attention = attention
        
        self.embedding = nn.Embedding(output_dim, emb_dim)
        
        self.rnn = nn.GRU((enc_hid_dim * 2) + emb_dim, dec_hid_dim)
        
        self.out = nn.Linear((enc_hid_dim * 2) + dec_hid_dim + emb_dim, output_dim)
        
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, input, hidden, encoder_outputs):
             
        #input = [batch size]
        #hidden = [batch size, dec hid dim]
        #encoder_outputs = [src sent len, batch size, enc hid dim * 2]
        
        input = input.unsqueeze(0)
        
        #input = [1, batch size]
        
        embedded = self.dropout(self.embedding(input))
        
        #embedded = [1, batch size, emb dim]
        
        a = self.attention(hidden, encoder_outputs)
                
        #a = [batch size, src len]
        
        a = a.unsqueeze(1)
        
        #a = [batch size, 1, src len]
        
        encoder_outputs = encoder_outputs.permute(1, 0, 2)
        
        #encoder_outputs = [batch size, src sent len, enc hid dim * 2]
        
        weighted = torch.bmm(a, encoder_outputs)
        
        #weighted = [batch size, 1, enc hid dim * 2]
        
        weighted = weighted.permute(1, 0, 2)
        
        #weighted = [1, batch size, enc hid dim * 2]
        
        rnn_input = torch.cat((embedded, weighted), dim=2)
        
        #rnn_input = [1, batch size, (enc hid dim * 2) + emb dim]
            
        output, hidden = self.rnn(rnn_input, hidden.unsqueeze(0))
        
        #output = [sent len, batch size, dec hid dim * n directions]
        #hidden = [n layers * n directions, batch size, dec hid dim]
        
        #sent len, n layers and n directions will always be 1 in this decoder, therefore:
        #output = [1, batch size, dec hid dim]
        #hidden = [1, batch size, dec hid dim]
        #this also means that output == hidden
        assert (output == hidden).all()
        
        embedded = embedded.squeeze(0)
        output = output.squeeze(0)
        weighted = weighted.squeeze(0)
        
        output = self.out(torch.cat((output, weighted, embedded), dim=1))
        
        #output = [bsz, output dim]
        
        return output, hidden.squeeze(0)

Seq2Seq

这是第一个我们不必使编码器RNN和解码器RNN具有相同隐藏层维度的模型，但编码器必须是双向的。可以通过更改所有出现的内容enc_dim * 2来删除此要求，enc_dim * 2 if encoder_is_bidirectional else enc_dim。

这个seq2seq封装器与之前两个模型类似。唯一的区别是encoder返回最终隐藏状态（这是通过线性层的前向和后向编码器RNN的最终隐藏状态），作为编码器的初始隐藏状态；以及编码器的每个隐藏状态（前后隐藏状态堆叠在彼此之上）。我们还需要确保将hidden和encoder_outputs传递给解码器。

简要介绍所有步骤：

创建outputs张量以保存所有的预测结果，$ \hat{Y} $
源序列$ X $被送入编码器，输出$ z $和$ H $
解码器的初始隐藏状态设置为上下文向量，$ s_0 = z = h_T $
我们使用一个batch的<sos>token作为解码器的第一个input，$ y_1 $
然后我们在循环中解码：
- 把输入值$ y_t $，上一步的隐藏状态 $ s_ {t-1} $ 和编码器输出的$ H $传递进解码器
- 接收预测，$ \hat{y}_{t + 1} $，以及更新的隐藏状态，$ s_t $
- 然后我们决定是否要用teacher forcing，确定下一步的输入$ y_{t+1}$

class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder, device):
        super().__init__()
        
        self.encoder = encoder
        self.decoder = decoder
        self.device = device
        
    def forward(self, src, trg, teacher_forcing_ratio=0.5):
        
        #src = [src sent len, batch size]
        #trg = [trg sent len, batch size]
        #teacher_forcing_ratio is probability to use teacher forcing
        #e.g. if teacher_forcing_ratio is 0.75 we use teacher forcing 75% of the time
        
        batch_size = src.shape[1]
        max_len = trg.shape[0]
        trg_vocab_size = self.decoder.output_dim
        
        #tensor to store decoder outputs
        outputs = torch.zeros(max_len, batch_size, trg_vocab_size).to(self.device)
        
        #encoder_outputs is all hidden states of the input sequence, back and forwards
        #hidden is the final forward and backward hidden states, passed through a linear layer
        encoder_outputs, hidden = self.encoder(src)
                
        #first input to the decoder is the <sos> tokens
        output = trg[0,:]
        
        for t in range(1, max_len):
            output, hidden = self.decoder(output, hidden, encoder_outputs)
            outputs[t] = output
            teacher_force = random.random() < teacher_forcing_ratio
            top1 = output.max(1)[1]
            output = (trg[t] if teacher_force else top1)

        return outputs

Training the Seq2Seq Model

INPUT_DIM = len(SRC.vocab)
OUTPUT_DIM = len(TRG.vocab)
ENC_EMB_DIM = 256
DEC_EMB_DIM = 256
ENC_HID_DIM = 512
DEC_HID_DIM = 512
ENC_DROPOUT = 0.5
DEC_DROPOUT = 0.5

attn = Attention(ENC_HID_DIM, DEC_HID_DIM)
enc = Encoder(INPUT_DIM, ENC_EMB_DIM, ENC_HID_DIM, DEC_HID_DIM, ENC_DROPOUT)
dec = Decoder(OUTPUT_DIM, DEC_EMB_DIM, ENC_HID_DIM, DEC_HID_DIM, DEC_DROPOUT, attn)

model = Seq2Seq(enc, dec, device).to(device)

optimizer = optim.Adam(model.parameters())
pad_idx = TRG.vocab.stoi['<pad>']
criterion = nn.CrossEntropyLoss(ignore_index=pad_idx)

# define train model
def train(model, iterator, optimizer, criterion, clip):
    model.train()
    epoch_loss = 0
    for i, batch in enumerate(iterator):
        src = batch.src
        trg = batch.trg
        
        optimizer.zero_grad()
        output = model(src, trg)
        
        #trg = [trg sent len, batch size]
        #output = [trg sent len, batch size, output dim]
        
        output = output[1:].view(-1, output.shape[-1])
        trg = trg[1:].view(-1)
        
        #trg = [(trg sent len - 1) * batch size]
        #output = [(trg sent len - 1) * batch size, output dim]
        
        loss = criterion(output, trg)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
        optimizer.step()
        
        epoch_loss += loss.item()
        
    return epoch_loss / len(iterator)
    
def evaluate(model, iterator, criterion): 
    model.eval()
    epoch_loss = 0
    with torch.no_grad():
        for i, batch in enumerate(iterator):
            src = batch.src
            trg = batch.trg
            output = model(src, trg, 0) #turn off teacher forcing

            #trg = [trg sent len, batch size]
            #output = [trg sent len, batch size, output dim]

            output = output[1:].view(-1, output.shape[-1])
            trg = trg[1:].view(-1)

            #trg = [(trg sent len - 1) * batch size]
            #output = [(trg sent len - 1) * batch size, output dim]

            loss = criterion(output, trg)
            epoch_loss += loss.item()
        
    return epoch_loss / len(iterator)
    
def epoch_time(start_time, end_time):
    elapsed_time = end_time - start_time
    elapsed_mins = int(elapsed_time / 60)
    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
    return elapsed_mins, elapsed_secs
    
N_EPOCHS = 10
CLIP = 1
SAVE_DIR = 'models'
MODEL_SAVE_PATH = os.path.join(SAVE_DIR, 'tut3_model.pt')

best_valid_loss = float('inf')

if not os.path.isdir(f'{SAVE_DIR}'):
    os.makedirs(f'{SAVE_DIR}')

for epoch in range(N_EPOCHS):
    
    start_time = time.time()
    
    train_loss = train(model, train_iterator, optimizer, criterion, CLIP)
    valid_loss = evaluate(model, valid_iterator, criterion)
    
    end_time = time.time()
    
    epoch_mins, epoch_secs = epoch_time(start_time, end_time)
    
    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        torch.save(model.state_dict(), MODEL_SAVE_PATH)
    
    print(f'| Epoch: {epoch+1:03} | Time: {epoch_mins}m {epoch_secs}s| Train Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f} | Val. Loss: {valid_loss:.3f} | Val. PPL: {math.exp(valid_loss):7.3f} |')

在下一篇博文中，我们将使用相同的架构，但将使用适用于所有RNN架构的一些技巧——预先训练的嵌入(pre-trained embeddings)，打包填充序列(packed padded sequences)和屏蔽(mask)。 我们还将实现代码，这将允许我们在解码时查看RNN在会注意输入语句中哪些字。

seq2seq-（3）

双向GRU和注意力机制