基于注意力的机器翻译-CSDN博客

本文链接：https://blog.csdn.net/m0_71684530/article/details/139911731

本博客为安徽大学自然语言处理实验课课程笔记。

1.概述

本次实验分别实现了基于注意力机制的编码器—解码器的法语翻译成英语模型和基于Transfomer的日文翻译成中文模型。

2.基于含注意力机制的编码器—解码器实现将简短的法语翻译成英语

2.1数据的读取和预处理

import collections
import os
import io
import math
import torch
from torch import nn
import torch.nn.functional as F
import torchtext.vocab as Vocab
import torch.utils.data as Data
import sys
import d2lzh_pytorch as d2l

定义必要的特殊符号：

<pad>：作为填充，使每个序列等长

<bos>：表示序列的开始

<eos>：表示序列的结束

# 定义填充、序列开始和序列结束的符号
PAD, BOS, EOS = '<pad>', '<bos>', '<eos>'

# 设置环境变量，指定使用哪个GPU（如果有的话）
# "0" 表示使用第一个GPU设备，如果有多GPU，可以指定其他数字
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

# 检查是否有可用的GPU，如果有则使用GPU，否则使用CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

定义process_one_seq函数，用于对当前词进行填充并添加结束标记，然后将其添加到列表中

# 定义一个函数来处理一个序列，将其词记录在all_tokens中，并在序列后面添加PAD直到序列长度达到max_seq_len
def process_one_seq(seq_tokens, all_tokens, all_seqs, max_seq_len):
    # 将当前序列中的词添加到所有词的列表中
    all_tokens.extend(seq_tokens)
    # 在序列末尾添加EOS标记，然后使用PAD填充直到序列长度达到max_seq_len
    seq_tokens += [EOS] + [PAD] * (max_seq_len - len(seq_tokens) - 1)
    # 将处理后的序列添加到所有序列的列表中
    all_seqs.append(seq_tokens)

定义build_data函数用于构建词汇表，并将所有的词序列转换为词索引序列，最后将这些索引转换为PyTorch张量，以便用于后续的模型训练

# 定义一个函数来构建词典和序列数据
def build_data(all_tokens, all_seqs):
    # 使用所有词的计数来创建一个词汇表对象，同时指定特殊标记PAD, BOS, EOS
    vocab = Vocab.Vocab(collections.Counter(all_tokens),
                        specials=[PAD, BOS, EOS])
    # 将所有序列中的词转换为对应的词索引
    indices = [[vocab.stoi[w] for w in seq] for seq in all_seqs]
    # 将索引列表转换为PyTorch张量并返回词汇表和张量
    return vocab, torch.tensor(indices)

fr-en-small.txt保存了一个很小的法语-英语数据集，在这个数据集里，每一行是一对法语句子和它对应的英语句子，中间使用'\t'隔开。在read_data函数中，我们为法语词和英语词分别创建词典，法语词的索引和英语词的索引相互独立。

# 定义一个函数来读取数据，并处理成最大长度为max_seq_len的序列
def read_data(max_seq_len):
    # 初始化输入和输出词的列表以及对应的序列列表
    in_tokens, out_tokens, in_seqs, out_seqs = [], [], [], []
    # 打开文件并读取所有行
    with io.open('fr-en-small.txt') as f:
        lines = f.readlines()
    # 遍历文件中的每一行
    for line in lines:
        # 将每行按照制表符分割为输入和输出序列
        in_seq, out_seq = line.rstrip().split('\t')
        # 将输入和输出序列进一步分割为词的列表
        in_seq_tokens, out_seq_tokens = in_seq.split(' '), out_seq.split(' ')
        # 如果加上EOS后长于max_seq_len，则忽略掉此样本
        if max(len(in_seq_tokens), len(out_seq_tokens)) > max_seq_len - 1:
            continue
        # 处理输入序列，将其词添加到输入词列表，并构造填充后的序列
        process_one_seq(in_seq_tokens, in_tokens, in_seqs, max_seq_len)
        # 处理输出序列，将其词添加到输出词列表，并构造填充后的序列
        process_one_seq(out_seq_tokens, out_tokens, out_seqs, max_seq_len)
    # 使用输入词列表构造输入词汇表和Tensor数据集
    in_vocab, in_data = build_data(in_tokens, in_seqs)
    # 使用输出词列表构造输出词汇表和Tensor数据集
    out_vocab, out_data = build_data(out_tokens, out_seqs)
    # 返回输入词汇表、输出词汇表以及构造的Tensor数据集
    return in_vocab, out_vocab, Data.TensorDataset(in_data, out_data)

2.2 构建含注意力机制的编码器—解码器

编码器—解码器架构主要由两个部分组成：编码器（encoder）和解码器（decoder）。编码器是一个循环神经网络，通常使用LSTM或GRU，负责将一个不定长的输入序列变换成一个定长的语义向量；解码器也是一个循环神经网络，负责根据语义向量生成指定的序列（这个过程称为解码）。

模型的训练主要需要考虑三个部分：编码器的输入、解码器的输入和解码器的输出。

未引入注意力机制的编码器—解码器将整个序列的信息都压缩进一个固定长度的向量中，然而一个向量无法完全表示整个序列的信息，特别是整个序列较长时。因而引入注意力机制，在每个时间输入不同的C（语义向量），使模型在不同位置的关注度不一样。

注意力机制通过对编码器所有时间步的隐状态做加权平均来得到语义向量。通过引入注意力机制，解码器输出序列的每个词条都会依赖一个与“上下文”相关的可变语义向量，而不再依赖一个相同的语义向量。

2.2.1 编码器

在编码器中，我们将输入语言的词索引通过词嵌入层得到词的表征，然后输入到一个多层门控循环单元中。

PyTorch的nn.GRU实例在前向计算后会分别返回输出和最终时间步的多层隐藏状态。其中的输出指的是最后一层的隐藏层在各个时间步的隐藏状态，并不涉及输出层计算。注意力机制将这些输出作为键项和值项。

# 定义编码器类，继承自nn.Module
class Encoder(nn.Module):
    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,
                 drop_prob=0, **kwargs):
        # 调用父类初始化方法
        super(Encoder, self).__init__(**kwargs)
        # 定义嵌入层，将词汇索引映射为嵌入向量
        self.embedding = nn.Embedding(vocab_size, embed_size)
        # 定义循环神经网络层，这里使用GRU
        self.rnn = nn.GRU(embed_size, num_hiddens, num_layers, dropout=drop_prob)

    def forward(self, inputs, state):
        # 输入形状是(批量大小, 时间步数)。将输出互换样本维和时间步维
        # 通过嵌入层获取嵌入向量
        embedding = self.embedding(inputs.long())
        # 将嵌入向量的维度从(批量大小, 时间步数, 嵌入维度)调整为(时间步数, 批量大小, 嵌入维度)
        embedding = embedding.permute(1, 0, 2)
        # 将处理后的嵌入向量和初始状态传递给GRU层
        return self.rnn(embedding, state)

    def begin_state(self):
        # 对于GRU，如果没有提供初始状态，PyTorch会自动初始化为零
        return None

2.2.2 注意力机制

定义attention_model函数，用于将输入连结后通过含单隐藏层的多层感知机变换。其中隐藏层的输入是解码器的隐藏状态与编码器在所有时间步上隐藏状态的连结，且使用tanh函数作为激活函数。输出层的输出个数为1。两个Linear实例均不使用偏差。其中函数a定义里向量v的长度是一个超参数。

def attention_model(input_size, attention_size):
    model = nn.Sequential(nn.Linear(input_size, attention_size, bias=False),
                          nn.Tanh(),
                          nn.Linear(attention_size, 1, bias=False))
    return model

注意力机制的输入包括查询项、键项和值项。设编码器和解码器的隐藏单元个数相同。这里的查询项为解码器在上一时间步的隐藏状态，形状为(批量大小, 隐藏单元个数)；键项和值项均为编码器在所有时间步的隐藏状态，形状为(时间步数, 批量大小, 隐藏单元个数)。注意力机制返回当前时间步的背景变量，形状为(批量大小, 隐藏单元个数)。

# 定义一个函数来计算注意力权重和背景向量
def attention_forward(model, enc_states, dec_state):
    """
    计算注意力权重和背景向量。

    参数:
    - model: 注意力模型，用于计算能量（注意力权重）
    - enc_states: 编码器的隐藏状态，形状为(时间步数, 批量大小, 隐藏单元个数)
    - dec_state: 解码器的当前隐藏状态，形状为(批量大小, 隐藏单元个数)

    返回:
    - context: 注意力加权的编码器隐藏状态，形状为(批量大小, 隐藏单元个数)
    """
    # 将解码器隐藏状态扩展到与编码器隐藏状态相同的时间步数
    dec_states = dec_state.unsqueeze(dim=0).expand_as(enc_states)
    # 将编码器和解码器的隐藏状态连结起来，以便计算注意力权重
    enc_and_dec_states = torch.cat((enc_states, dec_states), dim=2)
    # 使用注意力模型计算能量（注意力权重），形状为(时间步数, 批量大小, 1)
    e = model(enc_and_dec_states)
    # 在时间步维度对能量进行softmax运算，得到注意力权重
    alpha = F.softmax(e, dim=0)
    # 使用注意力权重对编码器隐藏状态进行加权，并求和得到背景变量
    context = (alpha * enc_states).sum(dim=0)
    # 返回背景变量，用于解码器的下一个时间步
    return context

2.2.3 含注意力机制的解码器

我们直接将编码器在最终时间步的隐藏状态作为解码器的初始隐藏状态。这要求编码器和解码器的循环神经网络使用相同的隐藏层个数和隐藏单元个数。

在解码器的前向计算中，我们先通过刚刚介绍的注意力机制计算得到当前时间步的背景向量。由于解码器的输入来自输出语言的词索引，我们将输入通过词嵌入层得到表征，然后和背景向量在特征维连结。我们将连结后的结果与上一时间步的隐藏状态通过门控循环单元计算出当前时间步的输出与隐藏状态。最后，我们将输出通过全连接层变换为有关各个输出词的预测，形状为(批量大小, 输出词典大小)。

# 定义解码器类，继承自nn.Module
class Decoder(nn.Module):
    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,
                 attention_size, drop_prob=0):
        super(Decoder, self).__init__()
        # 定义嵌入层，将词汇索引映射为嵌入向量
        self.embedding = nn.Embedding(vocab_size, embed_size)
        # 定义注意力机制模型
        self.attention = attention_model(2*num_hiddens, attention_size)
        # 定义循环神经网络层，这里使用GRU
        # GRU的输入尺寸是嵌入向量尺寸加上注意力机制的输出尺寸
        self.rnn = nn.GRU(embed_size + num_hiddens, num_hiddens, 
                          num_layers, dropout=drop_prob)
        # 定义输出层，将GRU的输出映射到词汇表大小
        self.out = nn.Linear(num_hiddens, vocab_size)

    def forward(self, cur_input, state, enc_states):
        """
        解码器的前向传播。

        参数:
        - cur_input: 当前输入，形状为(批量大小, )
        - state: 解码器的隐藏状态，形状为(num_layers, 批量大小, num_hiddens)
        - enc_states: 编码器的隐藏状态，形状为(时间步数, 批量大小, num_hiddens)

        返回:
        - output: 解码器的输出，形状为(批量大小, 词汇表大小)
        - state: 解码器的新状态，形状为(num_layers, 批量大小, num_hiddens)
        """
        # 使用注意力机制计算背景向量
        c = attention_forward(self.attention, enc_states, state[-1])
        # 将嵌入后的输入和背景向量在特征维连结
        input_and_c = torch.cat((self.embedding(cur_input), c), dim=1)
        # 为输入和背景向量的连结增加时间步维，时间步个数为1
        output, state = self.rnn(input_and_c.unsqueeze(0), state)
        # 移除时间步维，输出形状为(批量大小, 输出词典大小)
        output = self.out(output).squeeze(dim=0)
        return output, state

    def begin_state(self, enc_state):
        # 直接将编码器最终时间步的隐藏状态作为解码器的初始隐藏状态
        return enc_state

2.3 训练模型

batch_loss函数用于计算一个小批量的损失。解码器在最初时间步的输入是特殊BOS。之后，解码器在某时间步的输入为样本输出序列在上一时间步的词，即强制教学。为避免填充项对损失函数计算的影响，使用掩码变量。

def batch_loss(encoder, decoder, X, Y, loss):
    batch_size = X.shape[0]
    enc_state = encoder.begin_state()
    enc_outputs, enc_state = encoder(X, enc_state)
    # 初始化解码器的隐藏状态
    dec_state = decoder.begin_state(enc_state)
    # 解码器在最初时间步的输入是BOS
    dec_input = torch.tensor([out_vocab.stoi[BOS]] * batch_size)
    # 我们将使用掩码变量mask来忽略掉标签为填充项PAD的损失, 初始全1
    mask, num_not_pad_tokens = torch.ones(batch_size,), 0
    l = torch.tensor([0.0])
    for y in Y.permute(1,0): # Y shape: (batch, seq_len)
        dec_output, dec_state = decoder(dec_input, dec_state, enc_outputs)
        l = l + (mask * loss(dec_output, y)).sum()
        dec_input = y  # 使用强制教学
        num_not_pad_tokens += mask.sum().item()
        # EOS后面全是PAD. 下面一行保证一旦遇到EOS接下来的循环中mask就一直是0
        mask = mask * (y != out_vocab.stoi[EOS]).float()
    return l / num_not_pad_tokens

在训练函数中，我们需要同时迭代编码器和解码器的模型参数。

def train(encoder, decoder, dataset, lr, batch_size, num_epochs):
    enc_optimizer = torch.optim.Adam(encoder.parameters(), lr=lr)
    dec_optimizer = torch.optim.Adam(decoder.parameters(), lr=lr)

    loss = nn.CrossEntropyLoss(reduction='none')
    data_iter = Data.DataLoader(dataset, batch_size, shuffle=True)
    for epoch in range(num_epochs):
        l_sum = 0.0
        for X, Y in data_iter:
            enc_optimizer.zero_grad()
            dec_optimizer.zero_grad()
            l = batch_loss(encoder, decoder, X, Y, loss)
            l.backward()
            # 更新编码器参数
            enc_optimizer.step()
            # 更新解码器参数
            dec_optimizer.step()
            l_sum += l.item()
        if (epoch + 1) % 10 == 0:
            print("epoch %d, loss %.3f" % (epoch + 1, l_sum / len(data_iter)))

创建模型实例并设置超参数.

embed_size, num_hiddens, num_layers = 64, 64, 2
attention_size, drop_prob, lr, batch_size, num_epochs = 10, 0.5, 0.01, 2, 50
encoder = Encoder(len(in_vocab), embed_size, num_hiddens, num_layers,
                  drop_prob)
decoder = Decoder(len(out_vocab), embed_size, num_hiddens, num_layers,
                  attention_size, drop_prob)
train(encoder, decoder, dataset, lr, batch_size, num_epochs)

2.4 使用贪婪搜索生成解码器在每个时间步的输出

def translate(encoder, decoder, input_seq, max_seq_len):
    # 将输入序列分割为词列表，并在末尾添加EOS和PAD，确保长度为max_seq_len
    in_tokens = input_seq.split(' ')
    in_tokens += [EOS] + [PAD] * (max_seq_len - len(in_tokens) - 1)
    # 将输入词转换为词索引，并添加批次维，因为PyTorch期望输入至少有两维
    enc_input = torch.tensor([[in_vocab.stoi[tk] for tk in in_tokens]]) # batch=1
    # 获取编码器的初始状态
    enc_state = encoder.begin_state()
    # 通过编码器前向传播，得到编码器的输出和最终状态
    enc_output, enc_state = encoder(enc_input, enc_state)
    # 初始化解码器的输入为BOS（序列开始）的词索引
    dec_input = torch.tensor([out_vocab.stoi[BOS]])
    # 获取解码器的初始状态，这里使用编码器的最终状态
    dec_state = decoder.begin_state(enc_state)
    # 初始化输出词列表
    output_tokens = []
    # 循环执行解码器的最大序列长度次
    for _ in range(max_seq_len):
        # 通过解码器前向传播，得到解码器的输出和新状态
        dec_output, dec_state = decoder(dec_input, dec_state, enc_output)
        # 获取概率最高的词的索引
        pred = dec_output.argmax(dim=1)
        # 将词索引转换为词
        pred_token = out_vocab.itos[int(pred.item())]
        # 如果预测的词是EOS，则结束循环
        if pred_token == EOS:
            break
        else:
            # 否则，将预测的词添加到输出词列表
            output_tokens.append(pred_token)
            # 将预测的词作为下一次解码的输入
            dec_input = pred
    # 返回输出词列表
    return output_tokens

2.5 评价翻译结果

使用BLEU评价翻译结果，对于模型预测序列中任意的子序列，BLEU考察这个子序列是否出现在标签序列中。设词数为 $n$ 的子序列的精度为 $P_n$ 。它是预测序列与标签序列匹配词数为 $n$ 的子序列的数量与预测序列中词数为 $n$ 的子序列的数量之比。

例如：设标签序列为𝐴、𝐵、𝐶、𝐷、𝐸、𝐹，预测序列为𝐴、𝐵、𝐵、𝐶、𝐷，那么 $p_1=4/5, p_2=3/4, p_3=1/3, p_4=0$ ，设标签序列的词数为： $\text{len}_{\text{label}}$ ；预测序列的词数为： $\text{len}_{\text{pred}}$ 。则BLEU的定义为:

def bleu(pred_tokens, label_tokens, k):
    len_pred, len_label = len(pred_tokens), len(label_tokens)
    score = math.exp(min(0, 1 - len_label / len_pred))
    for n in range(1, k + 1):
        num_matches, label_subs = 0, collections.defaultdict(int)
        # 统计标签序列中所有可能的n-gram及其出现次数
        for i in range(len_label - n + 1):
            label_subs[''.join(label_tokens[i: i + n])] += 1
        # 遍历预测序列中的n-gram，并与标签序列中的n-gram进行匹配
        for i in range(len_pred - n + 1):
            if label_subs[''.join(pred_tokens[i: i + n])] > 0:
                num_matches += 1
                label_subs[''.join(pred_tokens[i: i + n])] -= 1
        # 使用0.5的n次幂作为权重，给予更低的n-gram更高的权重
        score *= math.pow(num_matches / (len_pred - n + 1), math.pow(0.5, n))
    return score


def score(input_seq, label_seq, k):
    pred_tokens = translate(encoder, decoder, input_seq, max_seq_len)
    label_tokens = label_seq.split(' ')
    print('bleu %.3f, predict: %s' % (bleu(pred_tokens, label_tokens, k),
                                      ' '.join(pred_tokens)))

3.基于Tansfomer的日文—中文翻译

import math
import torchtext
import torch
import torch.nn as nn
from torch import Tensor
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader
from collections import Counter
from torchtext.vocab import Vocab
from torch.nn import TransformerEncoder, TransformerDecoder, TransformerEncoderLayer, TransformerDecoderLayer
import io
import time
import pandas as pd
import numpy as np
import pickle
import tqdm
import sentencepiece as spm
torch.manual_seed(0)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

3.1 数据的读取和预处理

# 使用pandas读取平行语料库文件，分隔符为制表符，引擎设置为python，不使用标题行
df = pd.read_csv('zh-ja.bicleaner05.txt', sep='\t', engine='python', header=None)

# 将数据集中的第三列（汉语句子）转换为列表
trainen = df[2].values.tolist()

# 将数据集中的第四列（日语句子）转换为列表
trainja = df[3].values.tolist()

由于汉语和日语与英语不同，它们不包含空格来分隔单词，因而我们需要使用第三方分词器进行分词。

def build_vocab(sentences, tokenizer):
    counter = Counter()
    for sentence in sentences:
        # 使用分词器对句子进行编码，得到词的列表，并将词转换为字符串类型
        counter.update(tokenizer.encode(sentence, out_type=str))
    # 使用Counter对象和特殊标记列表创建一个Vocab对象
    return Vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>'])

# 使用日语句子列表trainja和日语分词器ja_tokenizer构建日语词汇表
ja_vocab = build_vocab(trainja, ja_tokenizer)
# 使用汉语句子列表trainen和汉语分词器en_tokenizer构建汉语词汇表
en_vocab = build_vocab(trainen, en_tokenizer)

构建训练数据的张量

def data_process(ja, en):
  data = []
  for (raw_ja, raw_en) in zip(ja, en):
    # 使用ja_tokenizer对日语句子进行编码，并去除句子末尾的换行符
    # 将编码后的单词转换为ja_vocab中的索引，并创建一个torch长整型张量
    ja_tensor_ = torch.tensor([ja_vocab[token] for token in ja_tokenizer.encode(raw_ja.rstrip("\n"), out_type=str)],
                            dtype=torch.long)
    # 对汉语句子进行相同的处理
    en_tensor_ = torch.tensor([en_vocab[token] for token in en_tokenizer.encode(raw_en.rstrip("\n"), out_type=str)],
                            dtype=torch.long)
    data.append((ja_tensor_, en_tensor_))
  return data
train_data = data_process(trainja, trainen)

为张量添加开始和结束标记，并填充到相同长度。

BATCH_SIZE = 8
PAD_IDX = ja_vocab['<pad>']
BOS_IDX = ja_vocab['<bos>']
EOS_IDX = ja_vocab['<eos>']
def generate_batch(data_batch):
  ja_batch, en_batch = [], []
  for (ja_item, en_item) in data_batch:
    # 在日语句子和汉语句子前后分别添加开始和结束标记，并连接成一个张量
    ja_batch.append(torch.cat([torch.tensor([BOS_IDX]), ja_item, torch.tensor([EOS_IDX])], dim=0))
    en_batch.append(torch.cat([torch.tensor([BOS_IDX]), en_item, torch.tensor([EOS_IDX])], dim=0))
  # 将日语和汉语批次中的句子填充到相同长度
  ja_batch = pad_sequence(ja_batch, padding_value=PAD_IDX)
  en_batch = pad_sequence(en_batch, padding_value=PAD_IDX)
  return ja_batch, en_batch
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=generate_batch)

3.2 基于Transfomer的seq2seq模型

模型搭建：

from torch.nn import (TransformerEncoder, TransformerDecoder,
                      TransformerEncoderLayer, TransformerDecoderLayer)


class Seq2SeqTransformer(nn.Module):
    def __init__(self, num_encoder_layers: int, num_decoder_layers: int,
                 emb_size: int, src_vocab_size: int, tgt_vocab_size: int,
                 dim_feedforward:int = 512, dropout:float = 0.1):
        super(Seq2SeqTransformer, self).__init__()
        
        encoder_layer = TransformerEncoderLayer(d_model=emb_size, nhead=NHEAD,
                                                dim_feedforward=dim_feedforward)
        # 创建Transformer编码器，它由num_encoder_layers个编码器层组成
        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        
        decoder_layer = TransformerDecoderLayer(d_model=emb_size, nhead=NHEAD,
                                                dim_feedforward=dim_feedforward)
        # 创建Transformer解码器，它由num_decoder_layers个解码器层组成
        self.transformer_decoder = TransformerDecoder(decoder_layer, num_layers=num_decoder_layers)

        # 定义输出层的全连接层，它将嵌入大小映射到目标词汇的大小
        self.generator = nn.Linear(emb_size, tgt_vocab_size)
        # 定义源语言的词嵌入层
        self.src_tok_emb = TokenEmbedding(src_vocab_size, emb_size)
        # 定义目标语言的词嵌入层
        self.tgt_tok_emb = TokenEmbedding(tgt_vocab_size, emb_size)
        # 定义位置编码层
        self.positional_encoding = PositionalEncoding(emb_size, dropout=dropout)
    # 定义前向传播函数    
    def forward(self, src: Tensor, trg: Tensor, src_mask: Tensor,
                tgt_mask: Tensor, src_padding_mask: Tensor,
                tgt_padding_mask: Tensor, memory_key_padding_mask: Tensor):
        # 对源语言进行词嵌入和位置编码
        src_emb = self.positional_encoding(self.src_tok_emb(src))
        # 对目标语言进行词嵌入和位置编码
        tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg))
        # 通过编码器处理源语言嵌入
        memory = self.transformer_encoder(src_emb, src_mask, src_padding_mask)
        # 通过解码器处理目标语言嵌入和编码器输出
        outs = self.transformer_decoder(tgt_emb, memory, tgt_mask, None,
                                        tgt_padding_mask, memory_key_padding_mask)
        return self.generator(outs)

    def encode(self, src: Tensor, src_mask: Tensor):
        return self.transformer_encoder(self.positional_encoding(
                            self.src_tok_emb(src)), src_mask)

    def decode(self, tgt: Tensor, memory: Tensor, tgt_mask: Tensor):
        return self.transformer_decoder(self.positional_encoding(
                          self.tgt_tok_emb(tgt)), memory,
                          tgt_mask)

使用标记嵌入表示文本标记。为了引入词序的概念，向标记嵌入中添加了位置编码。

class PositionalEncoding(nn.Module):
    def __init__(self, emb_size: int, dropout, maxlen: int = 5000):
        super(PositionalEncoding, self).__init__()
        den = torch.exp(- torch.arange(0, emb_size, 2) * math.log(10000) / emb_size)
        pos = torch.arange(0, maxlen).reshape(maxlen, 1)
        pos_embedding = torch.zeros((maxlen, emb_size))
        # 填充位置编码矩阵的偶数列（sin部分）
        pos_embedding[:, 0::2] = torch.sin(pos * den)
        # 填充位置编码矩阵的奇数列（cos部分）
        pos_embedding[:, 1::2] = torch.cos(pos * den)
        pos_embedding = pos_embedding.unsqueeze(-2)

        self.dropout = nn.Dropout(dropout)
        self.register_buffer('pos_embedding', pos_embedding)

    def forward(self, token_embedding: Tensor):
        # 将位置编码与词嵌入相加，并应用dropout
        return self.dropout(token_embedding +
                            self.pos_embedding[:token_embedding.size(0),:])

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size: int, emb_size):
        super(TokenEmbedding, self).__init__()
        self.embedding = nn.Embedding(vocab_size, emb_size)
        self.emb_size = emb_size
    def forward(self, tokens: Tensor):
        # 对输入的tokens进行词嵌入，并乘以根号下的嵌入大小进行缩放
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

使用generate_square_subsequent_mask函数创建方形后续掩码，作用于解码器，防止它关注序列中它之后的单词。

创建create_mask函数，作用于解码器和编码器，用于防止解码器和编码器关注未来的信息，并屏蔽填充的单词。

def generate_square_subsequent_mask(sz):
    # 创建一个大小为sz的正方形矩阵，其对角线以上的元素为1，以下的元素为0
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    # 将上三角矩阵中的1转换为0，0转换为负无穷大，得到一个掩码矩阵
    mask = mask.float().masked_fill(mask == 1, float('-inf')).masked_fill(mask == 0, float(0.0))
    return mask

def create_mask(src, tgt):
  src_seq_len = src.shape[0]
  tgt_seq_len = tgt.shape[0]

  # 生成目标语言序列的方形后续掩码
  tgt_mask = generate_square_subsequent_mask(tgt_seq_len)
  # 创建源语言序列的掩码，由于源语言序列通常不需要后续掩码，所以这里创建一个全0的布尔掩码
  src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool)

  src_padding_mask = (src == PAD_IDX).transpose(0, 1)
  tgt_padding_mask = (tgt == PAD_IDX).transpose(0, 1)
  return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask

模型的训练与评估：

SRC_VOCAB_SIZE = len(ja_vocab)
TGT_VOCAB_SIZE = len(en_vocab)
EMB_SIZE = 512
NHEAD = 8
FFN_HID_DIM = 512
BATCH_SIZE = 20
NUM_ENCODER_LAYERS = 3
NUM_DECODER_LAYERS = 3
NUM_EPOCHS = 16
transformer = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS,
                                 EMB_SIZE, SRC_VOCAB_SIZE, TGT_VOCAB_SIZE,
                                 FFN_HID_DIM)

for p in transformer.parameters():
    if p.dim() > 1:
        nn.init.xavier_uniform_(p)

transformer = transformer.to(device)

loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)

optimizer = torch.optim.Adam(
    transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9
)
def train_epoch(model, train_iter, optimizer):
  model.train()
  losses = 0
  for idx, (src, tgt) in  enumerate(train_iter):
      src = src.to(device)
      tgt = tgt.to(device)

      tgt_input = tgt[:-1, :]

      src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)

      logits = model(src, tgt_input, src_mask, tgt_mask,
                                src_padding_mask, tgt_padding_mask, src_padding_mask)

      optimizer.zero_grad()

      tgt_out = tgt[1:,:]
      loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
      loss.backward()

      optimizer.step()
      losses += loss.item()
  return losses / len(train_iter)


def evaluate(model, val_iter):
  model.eval()
  losses = 0
  for idx, (src, tgt) in (enumerate(valid_iter)):
    src = src.to(device)
    tgt = tgt.to(device)

    tgt_input = tgt[:-1, :]

    src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)

    logits = model(src, tgt_input, src_mask, tgt_mask,
                              src_padding_mask, tgt_padding_mask, src_padding_mask)
    tgt_out = tgt[1:,:]
    loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
    losses += loss.item()
  return losses / len(val_iter)

可以使用以下代码开始训练。

for epoch in tqdm.tqdm(range(1, NUM_EPOCHS+1)):
  start_time = time.time()
  train_loss = train_epoch(transformer, train_iter, optimizer)
  end_time = time.time()
  print((f"Epoch: {epoch}, Train loss: {train_loss:.3f}, "
          f"Epoch time = {(end_time - start_time):.3f}s"))

3.3 使用贪婪搜索输出翻译结果

# 定义一个函数greedy_decode，用于进行贪心解码
def greedy_decode(model, src, src_mask, max_len, start_symbol):
    src = src.to(device)
    src_mask = src_mask.to(device)
    memory = model.encode(src, src_mask)
    ys = torch.ones(1, 1).fill_(start_symbol).type(torch.long).to(device)
    # 进行解码，直到生成最大长度或遇到结束标记
    for i in range(max_len-1):
        memory = memory.to(device)
        memory_mask = torch.zeros(ys.shape[0], memory.shape[0]).to(device).type(torch.bool)
        tgt_mask = (generate_square_subsequent_mask(ys.size(0))
                                    .type(torch.bool)).to(device)
        out = model.decode(ys, memory, tgt_mask)
        out = out.transpose(0, 1)
        prob = model.generator(out[:, -1])
        _, next_word = torch.max(prob, dim = 1)
        next_word = next_word.item()
        ys = torch.cat([ys,
                        torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=0)
        if next_word == EOS_IDX:
          break
    return ys


# 定义一个函数translate，用于进行翻译
def translate(model, src, src_vocab, tgt_vocab, src_tokenizer):
    model.eval()
    tokens = [BOS_IDX] + [src_vocab.stoi[tok] for tok in src_tokenizer.encode(src, out_type=str)]+ [EOS_IDX]
    num_tokens = len(tokens)
    src = (torch.LongTensor(tokens).reshape(num_tokens, 1) )
    src_mask = (torch.zeros(num_tokens, num_tokens)).type(torch.bool)
    # 使用greedy_decode进行贪心解码
    tgt_tokens = greedy_decode(model,  src, src_mask, max_len=num_tokens + 5, start_symbol=BOS_IDX).flatten()
    return " ".join([tgt_vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "")