PyTorch实战：机器翻译

发表于 | 分类于 ， ，

从WMT18下载机器翻译数据集 News Commentary v13，并从中取出中英部分的平行语料，该部分为两个以 .en 和 .zh 为后缀的文件，文件的每一行是一个中文或英文的句子，两者一一对应互为翻译。

数据预处理

我们准备构建一个简单的“中到英”机器翻译系统，所以我们把中文语料作为源语言，英文语料作为目标语言；我们使用结巴分词对中文进行分词，NLTK对英文进行分词处理。我们预处理要做的事情包括：生成源语言和目标语言的分词序列，构建词典，划分训练、验证与测试集。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
import nltk,jieba,pickle
from sklearn.model_selection import train_test_split

PAD_TOKEN = 0
UNK_TOKEN = 1
SOS_TOKEN = 2
EOS_TOKEN = 3

def get_word_index(vocab_dic, max_size, min_freq):
    vocab_list = sorted([_ for _ in vocab_dic.items() if _[1] >= min_freq],
                           key=lambda x: x[1], reverse=True)[:max_size]
    vocab_list.insert(PAD_TOKEN, ('<pad>', min_freq))
    vocab_list.insert(UNK_TOKEN, ('<unk>', min_freq))
    vocab_list.insert(SOS_TOKEN, ('<sos>', min_freq))
    vocab_list.insert(EOS_TOKEN, ('<eos>', min_freq))
    word2idx = {word_count[0]: idx for idx, word_count in enumerate(vocab_list)}
    idx2word = {idx: word_count[0] for idx, word_count in enumerate(vocab_list)}
    return word2idx,idx2word


def preprocessing(src_file, tgt_file, max_size=10000, min_freq=1):
    with open(src_file, encoding='utf-8') as fin:
        chinese = fin.read().strip().split('\n')
    with open(tgt_file, encoding='utf-8') as fin:
        english = fin.read().strip().split('\n')
    assert (len(chinese) == len(english))

    src_list = []
    tgt_list = []
    en_vocab = {}
    ch_vocab = {}
    for i in range(len(english)):
        print(i)
        ch_tokens = list(jieba.cut(chinese[i]))
        en_tokens = nltk.word_tokenize(english[i])
        for word in ch_tokens:
            ch_vocab[word] = ch_vocab.get(word, 0) + 1
        for word in en_tokens:
            word = word.lower()
            en_vocab[word] = en_vocab.get(word, 0) + 1

        src_list.append(ch_tokens)
        tgt_list.append(en_tokens)
  
    ch_word2idx, ch_idx2word = get_word_index(ch_vocab, max_size, min_freq)
    with open('./data/MT/ch_vocab.pkl', 'wb') as fout:
        pickle.dump((ch_word2idx, ch_idx2word), fout)
    print('Source lang vocab size is: {}'.format(len(ch_word2idx)))
    en_word2idx,en_idx2word = get_word_index(en_vocab, max_size, min_freq)
    with open('./data/MT/en_vocab.pkl', 'wb') as fout:
        pickle.dump((en_word2idx,en_idx2word), fout)
    print('Target lang vocab size is: {}'.format(len(en_word2idx)))

    # 划分训练、验证、测试集
    src_train, src_test, tgt_train, tgt_test = train_test_split(src_list, tgt_list, test_size=0.1, shuffle=True)
    src_train, src_eval, tgt_train, tgt_eval = train_test_split(src_train, tgt_train, test_size=0.1, shuffle=True)
    with open('./data/MT/train-eval-corpus.pkl', 'wb') as fout:
        pickle.dump((src_train, tgt_train, src_eval, tgt_eval), fout)
    with open('./data/MT/test-corpus.pkl', 'wb') as fout:
        pickle.dump((src_test, tgt_test), fout)
    print("train-corpus size is: {}, eval-corpus size is: {}, test-corpus size is: {}".format(
        len(src_train), len(src_eval), len(src_test)))

不同于分类等其他的任务，翻译需要处理的数据是源语言-目标语言这样的平行语料，为了触发翻译任务的开始与结束需要额外增加两个标记句首与句尾的特殊标识（SOS_TOKEN 与 EOS_TOKEN）。

数据加载

数据的读取与加载我们仍然使用 PyTorch 提供的 Dataset 与 DataLoader 来实现：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import torch
from torch.utils.data import Dataset,DataLoader
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

class MyDataset(Dataset):
    def __init__(self, src_word2idx, src_corpus, tgt_word2idx, tgt_corpus):
        self.src_data = []
        self.tgt_data = []
        for src_sentence in src_corpus:
            self.src_data.append([src_word2idx.get(w, UNK_TOKEN) for w in src_sentence])
        for tgt_sentence in tgt_corpus:
            tgt_tokens = [tgt_word2idx.get(w, UNK_TOKEN) for w in tgt_sentence]
            tgt_tokens.append(EOS_TOKEN)
            self.tgt_data.append(tgt_tokens)

    def __len__(self):
        return len(self.src_data)

    def __getitem__(self, idx):
        return self.src_data[idx],self.tgt_data[idx]

我们继承 Dataset 实现自己的数据读取，分别对源语言和目标语言进行单词编号的转换，同时在目标语言的最后增加一个句尾标识（EOS_TOKEN）。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
def create_loader(src_word2idx, src_corpus, tgt_word2idx, tgt_corpus, batch_size):
    dataset = MyDataset(src_word2idx, src_corpus, tgt_word2idx, tgt_corpus)

    def get_batch_padded(batch_data, idx, pad):
        # 以批量数据中的最长序列为基准补齐
        max_len = max([len(seq[idx]) for seq in batch_data])
        data = []
        seq_lengths = []
        for seq in batch_data:
            data.append(seq[idx] + [pad] * (max_len - len(seq[idx])))
            seq_lengths.append(len(seq[idx]))
        data = torch.LongTensor(data).to(device)
        return data, seq_lengths

    def collate_fn(batch_data):
        src_data, src_seq_lengths = get_batch_padded(batch_data, 0, PAD_TOKEN)
        tgt_data, tgt_seq_lengths = get_batch_padded(batch_data, 1, PAD_TOKEN)
        return src_data, src_seq_lengths, tgt_data, tgt_seq_lengths

    return DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True,
                      collate_fn=collate_fn, drop_last=True)

通过 DataLoader 来实现批量数据的生成，我们还是采取批量数据中最长数据的长度进行补齐，同时，我们需要记录批量数据中真实文本序列的长度（仅源语言有这样的需求），为了后面使用 PyTorch 提供的压缩 pad 操作。

定义网络

机器翻译是一个由源语言序列到目标语言序列的 Seq2Seq（Sequence to Sequence) 任务，对应的网络架构是编码器-解码器网络，该网络框架的核心思想是将编码器中的序列（源语言）转换成一个上下文向量 $C$，并将该向量带入到解码器的每一步解码过程中，即目标语言的生成。我们首先来定义编码器网络：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_packed_sequence, pack_padded_sequence

class EncoderRNN(nn.Module):
    def __init__(self, input_size, hidden_size, n_layers=1):
        super(EncoderRNN, self).__init__()
        self.n_layers = n_layers
        self.hidden_size = hidden_size
        self.embedding = nn.Embedding(input_size, hidden_size, padding_idx=PAD_TOKEN)
        self.gru = nn.GRU(hidden_size, hidden_size, batch_first=True,
                          num_layers=self.n_layers, bidirectional=True)

    def forward(self, input, input_lengths, hidden):
        # input: [b, sl]
        # hidden: [drection*n_layers, b, h]
        output = self.embedding(input)
        # output: [b, sl, h]
        output = pack_padded_sequence(output, input_lengths, batch_first=True, enforce_sorted=False)
        # output is a PackedSequence object
        output, hidden = self.gru(output, hidden)
        # output: [b, sl, h*direction]
        # hidden: [drection*n_layers, b, h]
        output, _ = pad_packed_sequence(output, batch_first=True, padding_value=PAD_TOKEN)
        output = output[:, :, :self.hidden_size] + output[:, :, self.hidden_size:]
        # output: [b, sl, h]
        return output, hidden

    def initHidden(self, batch_size):
        result = torch.zeros(self.n_layers*2, batch_size, self.hidden_size).to(device)
        return result

编码器网络包含一个 embedding 层，我们将 embedding 的维度也设置为 hidden_size 的大小，输入通过 embedding 层以后接入一个双向的 GRU 层，我们把 GRU 的 batch_first 设置为 True，方便数据的形状的前后兼容。我们约定一些符号，b 表示 batch_size，h 表示 hidden_size，sl 表示源语言序列的长度，tl 表示目标语言序列的长度，便于我们在代码中注释上每一层输入输出的张量形状变化。请重点关注这些形状的变化，加深对各种常用网络层的输入输出形状的认识。

我们配合使用了 pack_padded_sequence( ) 和 pad_paked_sequence( ) 用来在循环神经网络中更好地处理 padding 以后的一个 batch 的源语言序列，避免神经网络学习到过多 pad 标签的无意义信息。

如果只是把编码器的隐层作为一个上下文向量 $C$，带入到解码器的每一步解码，这种处理过于简单，相当于将编码器中的每个 token 无差别地对待。但我们可以想象在翻译任务中，当我们输出目标语言的某个 token 时，其实重点参考的是源语言中的某一个或某几个 token ，这种随着解码输出的不同重点参考不同编码的机制就是注意力机制。所谓 “重点参考”，其数学表现形式就是权重矩阵，生成权重矩阵的输入源则是来自 解码器的隐层输出 $h_t$ 和 编码器的隐层输出 $\bar{h}_s$，它们按照注意力分布的计算公式来计算：

注意力分配矩阵 $\text{A}$ 的元素 $\alpha_{ts}$ 表示 $h_t$ 收到 $\bar{h}_s$ 的注意力概率。常用的计算 attention 的 score 方法有三种：

  • 内积（dot）： $h_t^T\bar{h}_s$
  • 线性映射（general）： $h_t^TW_a\bar{h}_s$
  • 双线性映射（concat）： $v_a^T\text{tanh}(W_a[h_t;\bar{h}_s])$

理解了注意力机制的原理后，我们来动手实现：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
class Attention(nn.Module):
    def __init__(self, score_type, hidden_size):
        super(Attention, self).__init__()
        self.score_type = score_type
        self.hidden_size = hidden_size
        if score_type == 'general':
            self.attn = nn.Linear(hidden_size, hidden_size)
        elif score_type == 'concat':
            self.attn = nn.Linear(hidden_size * 2, hidden_size)
            self.v = nn.Parameter(torch.FloatTensor(1, hidden_size))

     
    def score(self, decoder_rnn_output, encoder_output):
        # decoder_rnn_ouput: [1, h]
        # encoder_rnn_ouput: [1, h]
        if self.score_type == 'dot':
            energy = decoder_rnn_output.squeeze(0).dot(encoder_output.squeeze(0))
        elif self.score_type == 'general':
            energy = self.attn(encoder_output)
            energy = decoder_rnn_output.squeeze(0).dot(energy.squeeze(0))
        elif self.score_type == 'concat':
            h_o = torch.cat((decoder_rnn_output, encoder_output), 1)
            energy = self.attn(h_o)
            energy = self.v.squeeze(0).dot(energy.squeeze(0))
        return energy

    # attn_weights: [b, tl, sl]
    def forward(self, rnn_outputs, encoder_outputs):
        # rnn_outputs: [b, tl, h]
        # encoder_outputs: [b, sl, h]
        tgt_seq_len = rnn_outputs.size()[1]
        src_seq_len = encoder_outputs.size()[1]
        batch_size = encoder_outputs.size()[0]

        if self.score_type == 'general':
            encoder_outputs = self.attn(encoder_outputs).transpose(1, 2)
            # encoder_outputs: [b, h, sl]
            attn_energies = rnn_outputs.bmm(encoder_outputs)
            # attn_energies: # [b, tl, sl] <- [b, tl, h]*[b, h, sl]
            attn_weights = F.softmax(attn_energies, dim=1)
            return attn_weights
        else:
            attn_energies = torch.zeros(batch_size, tgt_seq_len, src_seq_len)
            # attn_energies: [b, tl, sl]
            for b in range(batch_size):
                decoder_rnn_output = rnn_outputs[b]
                # decoder_rnn_ouput: [tl, h]
                for i in range(src_seq_len):
                    encoder_output = encoder_outputs[b, i, :].unsqueeze(0)
                    attn_energies[b, :, i] = self.score(decoder_rnn_output, encoder_output)

            attn_weights = torch.zeros(batch_size, src_seq_len).to(device)
            for b in range(batch_size):
                attn_weights[b] = F.softmax(attn_energies[b], dim=1)
            return attn_weights.unsqueeze(1)

以上的代码其实就是对上述的注意力原理与公式的翻译，请大家关注注释中各张量的形状变化，其实如果我们读到后文训练部分时再回头来看上面的代码，你会发现其实 tl 就等于 1。实现完了注意力层以后，就可以开始实现解码层了：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class AttnDecoderRNN(nn.Module):
    def __init__(self, hidden_size, output_size, score_method='concat', n_layers=1, dropout_p=0.1):
        super(AttnDecoderRNN, self).__init__()
        self.score_method = score_method
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.n_layers = n_layers
        self.dropout_p = dropout_p

        self.embedding = nn.Embedding(output_size, hidden_size, padding_idx=PAD_TOKEN)
        self.embedding_dropout = nn.Dropout(dropout_p)
        self.gru = nn.GRU(hidden_size, hidden_size, batch_first=True, n_layers=n_layers)
        self.attn = Attention(score_method, hidden_size)
        self.concat = nn.Linear(hidden_size * 2, hidden_size)
        self.out = nn.Linear(hidden_size, output_size)
        self.softmax = nn.LogSoftmax(dim=1)

    def forward(self, input_seqs, last_hidden, encoder_outputs):
        # input_seqs: [b, tl]
        # last_hidden: [direction*n_layers, b, h]
        # encoder_outputs: [b, sl, h]

        embedded = self.embedding(input_seqs)
        # embedded: [b, tl, h]
        rnn_output, hidden = self.gru(embedded, last_hidden)
        # rnn_outupt: [b, tl, h]
        # hidden: [direction*n_layers, b, h]
        attn_weights = self.attn(rnn_output, encoder_outputs)
        # attn_weights: [b, tl, sl]
        context = attn_weights.bmm(encoder_outputs)
        # context: [b, tl, h] <- [b, tl, sl] * [b, sl, h]
        output_context = torch.cat((rnn_output, context), 2)
        # output_context: [b, tl, 2h] <- [b, tl, h], [b, tl, h]
        output_context = self.concat(output_context)
        # output_context: [b, tl, h]
        concat_output = torch.tanh(output_context)
        
        output = self.out(concat_output)
        output = F.softmax(output, dim=1)

        return output, hidden

解码层与编码层的逻辑是差不多的，embedding 层过后接入 GRU 层，不同的是这里的 GRU 是一个单向的；随后进行注意力计算，将 GRU 的输出与注意力值进行融合（线性变化加非线性激活），最后进行 softmax 层映射到目标语言的词汇。

开始训练

训练过程我们引入了一种叫做 “教师强制（teacher forcing）” 的机制，即解码器在生成下一个词的时候，输入的不是自己预测的词，而是真实数据。通过设置 teacher_forcing_ratio=0.5 即代表着将有 50% 的概率选用这种直接使用正确答案的监督学习方式，余下的 50% 直接用预测的结果作为输入。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
import random

def train_simple_model(train_loader, encoder, decoder, encoder_optimizer, decoder_optimizer, criterion, config):

    total_loss = 0.0
    for src_input,src_lengths,tgt_input,tgt_lengths in train_loader:
        # 循环神经网络每一个 batch 都要重新初始化隐层
        encoder_hidden = encoder.initHidden(config.batch_size)
        encoder_optimizer.zero_grad()
        decoder_optimizer.zero_grad()

        loss = torch.tensor([0.0]).to(device)
        encoder_outputs, encoder_hidden = encoder(src_input, src_lengths, encoder_hidden)

        # 输入给解码器的第一个字符是句首标识 SOS_TOKEN
        tgt_seq_len = tgt_input.size()[1]
        decoder_input = torch.LongTensor([[SOS_TOKEN]] * config.batch_size).to(device)
        # 使用掩码变量mask来忽略掉标签为填充项PAD的损失
        mask, num_not_pad_tokens = torch.ones(config.batch_size, ).to(device), 0
        # 编码器是双向GRU，而解码器是单向的
        decoder_hidden = encoder_hidden[:decoder.n_layers]
        
        use_teacher_forcing = True if random.random() < config.teacher_forcing_ratio else False
        for t_step in range(tgt_seq_len):
            # 开始一步一步地解码
            decoder_output, decoder_hidden = decoder(decoder_input, decoder_hidden, encoder_outputs)
            loss = loss + (mask * criterion(decoder_output, tgt_input[:, t_step])).sum()
            num_not_pad_tokens += mask.sum().item()
            # EOS 后面全是 PAD，保证一旦遇到EOS接下来的循环中 mask 就一直是 0
            mask = mask * (decoder_input != EOS_TOKEN).float()

            if use_teacher_forcing:
                # 将真实数据当做下一时间步的输入
                decoder_input = tgt_input[:,t_step].unsqueeze(1)  # Teacher forcing
            else:
                # 从输出结果（概率的对数值）中选择出一个数值最大的单词作为输出放到了topi中
                topv, topi = decoder_output.data.topk(1, dim=1)
                ni = topi[:, 0]
                decoder_input = ni.unsqueeze(1).to(device)
                # decoder_input: [b, 1]

        # 开始反向传播
        loss.backward()
        # 开始梯度下降
        encoder_optimizer.step()
        decoder_optimizer.step()
        # 累加总误差
        total_loss += (loss / num_not_pad_tokens).item()

    # 返回训练时候的平均误差
    return total_loss / len(train_loader)

从上面的代码可以看到，编码器在训练过程中不需要关心其内部的时间步展开与循环，由编码器网络直接返回 encoder_outputs 和 encoder_hidden；解码器的训练则需要按时间步一步一步地喂数据，而且在开始解码时输入的第一个字符应该是句首 SOS_TOKEN。另外，我们在计算损失的时候，使用了掩码变量将句尾标识 EOS_TOKEN 往后的 pad 值都忽略掉。

启动训练，我们把相关的超参数封装到 Config 对象中，并且完成数据的读取与加载；然后初始化网络、初始化优化器、定义损失函数；随后进行多轮迭代，并在每轮迭代中进行验证集上的验证，记录并输出损失的变化；最后，保存模型。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
from torch import optim
import numpy as np

class Config(object):
    batch_size = 16
    num_epochs = 200
    hidden_size = 64
    n_layers = 1
    learning_rate = 0.0001
    src_vocab_size = 10004
    tgt_vocab_size = 10004
    teacher_forcing_ratio = 0.5


def run():
    plot_losses = []

    src_word2idx, _ = pickle.load(open('./data/MT/ch_vocab.pkl', 'rb'))
    tgt_word2idx, _ = pickle.load(open('./data/MT/en_vocab.pkl', 'rb'))
    (src_train, tgt_train, src_eval, tgt_eval) = pickle.load(open('./data/MT/train-eval-corpus.pkl', 'rb'))

    config = Config()

    encoder = EncoderRNN(config.src_vocab_size, config.hidden_size, n_layers=config.n_layers).to(device)
    decoder = AttnDecoderRNN(config.hidden_size, config.tgt_vocab_size, score_method='general',
                             n_layers=config.n_layers, dropout_p=0.1).to(device)

    # 为两个网络分别定义优化器
    encoder_optimizer = optim.Adam(encoder.parameters(), lr=config.learning_rate)
    decoder_optimizer = optim.Adam(decoder.parameters(), lr=config.learning_rate)

    # 定义损失函数
    criterion = nn.NLLLoss()

    # 开始多轮迭代训练
    for epoch in range(config.num_epochs):
        train_loader = create_loader(src_word2idx, src_train, tgt_word2idx, tgt_train, config.batch_size)
        train_loss = train_simple_model(train_loader, encoder, decoder, encoder_optimizer, decoder_optimizer, criterion, config)

        valid_loader = create_loader(src_word2idx, src_eval, tgt_word2idx, tgt_eval, config.batch_size)
        valid_loss = evaluation_simple_model(valid_loader, encoder, decoder, criterion, config, rights)

        # 打印并记录每一个Epoch的输出结果
        print('Epoch: %d%% Train-Loss: %.4f Eval-Loss: %.4f' % 
        (epoch * 1.0 / config.num_epochs * 100, train_loss, valid_loss))
        plot_losses.append([train_loss, valid_loss])
    
    torch.save(encoder, './model/NMT_encoder.pth')
    torch.save(decoder, './model/NMT_decoder.pth')

由于篇幅有限，我们不再给出 evaluation_simple_model( ) 的实现，其处理流程与训练函数基本一致（不需要引入教师强制策略）。

使用模型

使用我们训练好的模型，我们来进行一次中到英的翻译，值得注意的是我们的模型是针对 mini-batch 数据格式的，而一般翻译任务的输入是单条数据，我们需要进行 batch 的封装：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
def translate(encoder_model, decoder_mdoel, max_seq_len):
    encoder = torch.load(encoder_model, map_location=device)
    decoder = torch.load(decoder_mdoel, map_location=device)
    ch_word2idx, _ = pickle.load(open('./data/MT/ch_vocab.pkl', 'rb'))
    _, en_idx2word = pickle.load(open('./data/MT/en_vocab.pkl', 'rb'))


    ch_text = "美国的立场也起不到帮助。"
    ch_tokens = list(jieba.cut(ch_text))
    ch_tokens.append('<eos>')
    enc_input = torch.tensor([[ch_word2idx.get(w, UNK_TOKEN) for w in ch_tokens]]).to(device)
    encoder_hidden = encoder.initHidden(1)
    encoder_outputs, encoder_hidden  = encoder(enc_input, [len(ch_tokens)], encoder_hidden)
    decoder_input = torch.LongTensor([[SOS_TOKEN]]).to(device)
    decoder_hidden = encoder_hidden[:1]
    output_tokens = []
    for _ in range(max_seq_len):
        decoder_output, decoder_hidden = decoder(decoder_input, decoder_hidden, encoder_outputs)
        pred = decoder_output.argmax(dim=1)
        pred_token = en_idx2word.get(int(pred[0].item()), '<unk>')
        if pred_token == EOS_TOKEN:
            break
        else:
            output_tokens.append(pred_token)
            decoder_input = pred.unsqueeze(0)
    print(output_tokens)

我们需要加载保存好的编码器与解码器模型和词典数据，另外，我们需要指定翻译输出的最大长度。实践中，评价机器翻译结果通常使用 BLEU（Bilingual Evaluation Understudy），对于模型预测序列中任意的子序列，BLEU 考察这个子序列是否出现在标签序列中，在这里我们不做详细介绍，感兴趣的读者可以查阅相关资料。