从零实现带注意力机制的神经翻译模型：PyTorch实战指南

在自然语言处理领域，机器翻译一直是最具挑战性的任务之一。传统的序列到序列(Seq2Seq)模型在处理长句子时表现不佳，而注意力机制的引入彻底改变了这一局面。本文将带您从零开始实现一个基于PyTorch的英法翻译模型，重点讲解Bahdanau注意力的实现细节，并通过对比实验展示注意力机制的实际效果。

1. 项目准备与环境搭建

首先我们需要配置开发环境并准备数据集。推荐使用Python 3.8+和PyTorch 1.10+版本，这些版本在兼容性和性能方面都有良好表现。

pip install torch torchtext spacy
python -m spacy download en_core_web_sm
python -m spacy download fr_core_news_sm

我们将使用英法平行语料库作为训练数据。TorchText库提供了方便的接口来加载和处理这类数据集：

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

# 定义字段处理器
SRC = Field(tokenize="spacy", tokenizer_language="fr", init_token="<sos>", eos_token="<eos>", lower=True)
TRG = Field(tokenize="spacy", tokenizer_language="en", init_token="<sos>", eos_token="<eos>", lower=True)

# 加载数据集
train_data, valid_data, test_data = Multi30k.splits(exts=(".fr", ".en"), fields=(SRC, TRG))

# 构建词汇表
SRC.build_vocab(train_data, min_freq=2)
TRG.build_vocab(train_data, min_freq=2)

# 创建数据迭代器
BATCH_SIZE = 128
train_iterator, valid_iterator, test_iterator = BucketIterator.splits(
    (train_data, valid_data, test_data), batch_size=BATCH_SIZE, device=device)

2. 基础Seq2Seq模型实现

在引入注意力机制前，我们先实现一个基础的编码器-解码器架构。编码器将源语言句子转换为上下文向量，解码器则基于该向量生成目标语言句子。

编码器实现：

import torch.nn as nn

class Encoder(nn.Module):
    def __init__(self, input_dim, emb_dim, hid_dim, n_layers, dropout):
        super().__init__()
        self.embedding = nn.Embedding(input_dim, emb_dim)
        self.rnn = nn.GRU(emb_dim, hid_dim, n_layers, dropout=dropout)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, src):
        # src形状: [src_len, batch_size]
        embedded = self.dropout(self.embedding(src))
        # embedded形状: [src_len, batch_size, emb_dim]
        outputs, hidden = self.rnn(embedded)
        # outputs形状: [src_len, batch_size, hid_dim * n_directions]
        # hidden形状: [n_layers * n_directions, batch_size, hid_dim]
        return hidden

基础解码器实现：

class Decoder(nn.Module):
    def __init__(self, output_dim, emb_dim, hid_dim, n_layers, dropout):
        super().__init__()
        self.output_dim = output_dim
        self.embedding = nn.Embedding(output_dim, emb_dim)
        self.rnn = nn.GRU(emb_dim + hid_dim, hid_dim, n_layers, dropout=dropout)
        self.fc_out = nn.Linear(emb_dim + hid_dim * 2, output_dim)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, input, hidden, context):
        # input形状: [batch_size]
        # hidden形状: [n_layers, batch_size, hid_dim]
        # context形状: [batch_size, hid_dim]
        input = input.unsqueeze(0)
        embedded = self.dropout(self.embedding(input))
        # embedded形状: [1, batch_size, emb_dim]
        emb_con = torch.cat((embedded, context.unsqueeze(0)), dim=2)
        output, hidden = self.rnn(emb_con, hidden)
        # output形状: [1, batch_size, hid_dim]
        # hidden形状: [n_layers, batch_size, hid_dim]
        output = torch.cat((embedded.squeeze(0), hidden[-1], context), dim=1)
        prediction = self.fc_out(output)
        # prediction形状: [batch_size, output_dim]
        return prediction, hidden

3. 注意力机制的原理与实现

注意力机制的核心思想是让解码器在生成每个词时，能够动态地关注源句子中最相关的部分。Bahdanau注意力(又称加法注意力)通过计算解码器当前状态与所有编码器状态的相似度得分，然后对这些得分进行softmax归一化得到注意力权重。

注意力计算过程：

1. 解码器当前隐藏状态与所有编码器隐藏状态计算相似度得分
2. 对得分进行softmax归一化得到注意力权重
3. 计算编码器隐藏状态的加权和作为上下文向量
4. 将上下文向量与解码器输入拼接后送入RNN

class Attention(nn.Module):
    def __init__(self, hid_dim):
        super().__init__()
        self.attn = nn.Linear(hid_dim * 2, hid_dim)
        self.v = nn.Linear(hid_dim, 1, bias=False)
        
    def forward(self, hidden, encoder_outputs):
        # hidden形状: [n_layers, batch_size, hid_dim]
        # encoder_outputs形状: [src_len, batch_size, hid_dim]
        src_len = encoder_outputs.shape[0]
        hidden = hidden[-1].unsqueeze(1).repeat(1, src_len, 1)
        encoder_outputs = encoder_outputs.permute(1, 0, 2)
        energy = torch.tanh(self.attn(torch.cat((hidden, encoder_outputs), dim=2)))
        attention = self.v(energy).squeeze(2)
        return F.softmax(attention, dim=1)

带注意力的解码器实现：

class AttnDecoder(nn.Module):
    def __init__(self, output_dim, emb_dim, hid_dim, n_layers, dropout):
        super().__init__()
        self.output_dim = output_dim
        self.attention = Attention(hid_dim)
        self.embedding = nn.Embedding(output_dim, emb_dim)
        self.rnn = nn.GRU(emb_dim + hid_dim, hid_dim, n_layers, dropout=dropout)
        self.fc_out = nn.Linear(emb_dim + hid_dim * 2, output_dim)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, input, hidden, encoder_outputs):
        # input形状: [batch_size]
        # hidden形状: [n_layers, batch_size, hid_dim]
        # encoder_outputs形状: [src_len, batch_size, hid_dim]
        input = input.unsqueeze(0)
        embedded = self.dropout(self.embedding(input))
        # embedded形状: [1, batch_size, emb_dim]
        a = self.attention(hidden, encoder_outputs)
        # a形状: [batch_size, src_len]
        encoder_outputs = encoder_outputs.permute(1, 0, 2)
        weighted = torch.bmm(a.unsqueeze(1), encoder_outputs)
        # weighted形状: [batch_size, 1, hid_dim]
        rnn_input = torch.cat((embedded, weighted.permute(1, 0, 2)), dim=2)
        output, hidden = self.rnn(rnn_input, hidden)
        # output形状: [1, batch_size, hid_dim]
        embedded = embedded.squeeze(0)
        output = output.squeeze(0)
        weighted = weighted.squeeze(1)
        prediction = self.fc_out(torch.cat((output, weighted, embedded), dim=1))
        # prediction形状: [batch_size, output_dim]
        return prediction, hidden, a

4. 完整模型训练与评估

现在我们将编码器和带注意力的解码器组合成完整的Seq2Seq模型，并实现训练和评估流程。

完整模型实现：

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):
        batch_size = src.shape[1]
        trg_len = trg.shape[0]
        trg_vocab_size = self.decoder.output_dim
        outputs = torch.zeros(trg_len, batch_size, trg_vocab_size).to(self.device)
        encoder_outputs, hidden = self.encoder(src)
        input = trg[0, :]
        for t in range(1, trg_len):
            output, hidden, _ = self.decoder(input, hidden, encoder_outputs)
            outputs[t] = output
            teacher_force = random.random() < teacher_forcing_ratio
            top1 = output.argmax(1)
            input = trg[t] if teacher_force else top1
        return outputs

训练函数：

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)
        output_dim = output.shape[-1]
        output = output[1:].view(-1, output_dim)
        trg = trg[1:].view(-1)
        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)  # 关闭teacher forcing
            output_dim = output.shape[-1]
            output = output[1:].view(-1, output_dim)
            trg = trg[1:].view(-1)
            loss = criterion(output, trg)
            epoch_loss += loss.item()
    return epoch_loss / len(iterator)

5. 注意力机制效果分析与可视化

为了直观理解注意力机制的作用，我们可以将训练好的模型在测试集上进行评估，并可视化注意力权重。

注意力可视化函数：

def plot_attention(attention, source, target):
    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(111)
    cax = ax.matshow(attention, cmap='bone')
    fig.colorbar(cax)
    ax.set_xticklabels([''] + source, rotation=90)
    ax.set_yticklabels([''] + target)
    ax.xaxis.set_major_locator(ticker.MultipleLocator(1))
    ax.yaxis.set_major_locator(ticker.MultipleLocator(1))
    plt.show()

注意力分析示例：

def translate_sentence(sentence, src_field, trg_field, model, device, max_len=50):
    model.eval()
    tokens = [token.lower() for token in sentence]
    tokens = [src_field.init_token] + tokens + [src_field.eos_token]
    src_indexes = [src_field.vocab.stoi[token] for token in tokens]
    src_tensor = torch.LongTensor(src_indexes).unsqueeze(1).to(device)
    encoder_outputs, hidden = model.encoder(src_tensor)
    trg_indexes = [trg_field.vocab.stoi[trg_field.init_token]]
    attentions = torch.zeros(max_len, 1, len(src_indexes)).to(device)
    for i in range(1, max_len):
        trg_tensor = torch.LongTensor([trg_indexes[-1]]).to(device)
        with torch.no_grad():
            output, hidden, attention = model.decoder(trg_tensor, hidden, encoder_outputs)
        attentions[i] = attention
        pred_token = output.argmax(1).item()
        trg_indexes.append(pred_token)
        if pred_token == trg_field.vocab.stoi[trg_field.eos_token]:
            break
    trg_tokens = [trg_field.vocab.itos[i] for i in trg_indexes]
    return trg_tokens[1:], attentions[:len(trg_tokens)-1]

通过对比实验可以发现，引入注意力机制的模型在长句子翻译任务上表现显著优于基础Seq2Seq模型。特别是在处理主语-动词一致性、性别一致性和复杂句式结构时，注意力机制能够更准确地捕捉源语言和目标语言之间的对应关系。