pytorch 实现 Seq2Seq 机器翻译

本文主要介绍了用 pytorch 实现简单的 Seq2Seq 机器翻译任务，参考了 李宏毅老师的深度学习课程 的作业八，数据集也来自于此，视频课程可以在B站学习（https://www.bilibili.com/video/BV1JE411g7XF?p=53）。算法理论可以阅读论文"Sequence to Sequence Learning with Neural Networks"，也可以参考 我写的论文阅读笔记 。

import 需要使用的模块

import torch
import torch.nn as nn
import torch.utils.data as data
import torchsummary
from torchvision import datasets
import numpy as np
import sys
import os
import random
import re
import json
from nltk.translate.bleu_score import sentence_bleu

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

设置需要的配置参数

data_path = "./cmn-eng"          # 数据集的位置
store_model_path = "./ckpt"      # 储存模型的位置
max_output_len = 45              # 输出句子的最大长度
batch_size = 64                  # batch_size
emb_dim = 256                    # word embedding向量的维度
hid_dim = 512                    # RNN隐藏状态的维度
n_layers = 4                     # RNN的层数
dropout = 0.5                    # dropout的概率p
learning_rate = 0.0001           # 初始化学习率
teacher_forcing_ratio = 0.5      # 使用正解训练的概率
summary_steps = 6000             # 总训练batch数

从下载的数据文件中导入数据集和中英文字典

# 载入字典
def get_dictionary(root, language):
    with open(os.path.join(root, 'word2int_{}.json'.format(language)), "r") as f:
        word2int = json.load(f)
    with open(os.path.join(root, 'int2word_{}.json'.format(language)), "r") as f:
        int2word = json.load(f)
    print('{} vocab size: {}'.format(language, len(word2int)))
    return word2int, int2word, len(word2int)

word2int_cn, int2word_cn, cn_vocab_size = get_dictionary(data_path, 'cn') # 中文字典
word2int_en, int2word_en, en_vocab_size = get_dictionary(data_path, 'en') # 英文字典
vocab=[word2int_cn, int2word_cn, word2int_en, int2word_en]

# 载入数据 (training/validation/testing)
def load_data(root, set_name):
    data = []
    with open(os.path.join(root, '{}.txt'.format(set_name)), "r") as f:
        for line in f:
            data.append(line)
    print('{} dataset size: {}'.format(set_name,len(data)))
    
    return data

training_data=load_data(data_path, 'training')
val_data=load_data(data_path, 'validation')
testing_data=load_data(data_path, 'testing')

'''
打印结果
cn vocab size: 3805
en vocab size: 3922
training dataset size: 18000
validation dataset size: 500
testing dataset size: 2636
'''

从输出结果可以看到，中英文字典中的词汇量分别为3805、3922，训练集、验证集、测试集中的句子数量分别为18000、500、2636。

对数据集进行一些处理，也就是继承 torch.utils.data.Dataset 类，然后实现里面的 __len__ 以及 __getitem__ 方法

class EN2CNDataset(data.Dataset):
    def __init__(self, data, max_output_len, vocab):
        self.max_output_len = max_output_len
        self.word2int_cn, self.int2word_cn = vocab[0], vocab[1] # 中文字典
        self.word2int_en, self.int2word_en = vocab[2], vocab[3] # 英文字典
        self.data = data

        self.cn_vocab_size = len(self.word2int_cn)
        self.en_vocab_size = len(self.word2int_en)

    def seq_pad(self, label, pad_token):
        # 将不同长度的句子pad到相同长度，以便训练
        label = np.pad(label, (0, (self.max_output_len - label.shape[0])), mode='constant', constant_values=pad_token)
        return label

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

    def __getitem__(self, Index):
        # 将英文句子和中文句子分开
        sentences = self.data[Index]
        sentences = re.split('[\t\n]', sentences)
        sentences = list(filter(None, sentences))
        #print (sentences)
        assert len(sentences) == 2

        # 特殊token
        BOS = self.word2int_en['<BOS>']
        EOS = self.word2int_en['<EOS>']
        UNK = self.word2int_en['<UNK>']

        # 在句子开头添加'<BOS>'，结尾添加'<EOS>'，字典中没有的词标记为'<UNK>'
        en, cn = [BOS], [BOS]
        # 英文句子分词后转为字典索引向量
        sentence = re.split(' ', sentences[0])
        sentence = list(filter(None, sentence))
        for word in sentence:
            en.append(self.word2int_en.get(word, UNK))
        en.append(EOS)

        # 中文句子分词后转为字典索引向量
        # e.g. < BOS >, we, are, friends, < EOS > --> 1, 28, 29, 205, 2
        sentence = re.split(' ', sentences[1])
        sentence = list(filter(None, sentence))
        for word in sentence:
            cn.append(self.word2int_cn.get(word, UNK))
        cn.append(EOS)

        en, cn = np.asarray(en), np.asarray(cn)
        #if len(en)>30 or len(cn)>30:
        #    print(len(en),len(cn)) 

        # 用 '<PAD>' 将句子pad到相同长度
        en = self.seq_pad(en, self.word2int_en['<PAD>'])
        cn = self.seq_pad(cn, self.word2int_cn['<PAD>'])
        en, cn = torch.LongTensor(en), torch.LongTensor(cn)

        # return英文和中文句子的向量
        return en, cn

train_dataset = EN2CNDataset(training_data, max_output_len, vocab)
val_dataset = EN2CNDataset(val_data, max_output_len, vocab)
test_dataset = EN2CNDataset(testing_data, max_output_len, vocab)
train_loader = data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = data.DataLoader(val_dataset, batch_size=1)
test_loader = data.DataLoader(test_dataset, batch_size=1)

主要处理为：在每个句子的首尾加上开始和结束符号 ('<BOS>' & '<EOS>') ；用 '<UNK>' 符号代替句子中没在字典中出现的词，表示生词；将每个句子用 '<PAD>' 符号补充到相同长度 max_output_len ；将句子中的词转换为字典中的index，以index序列的形式存储；中英文序列分开存储。

定义 Encoder 和 Decoder 类

class Encoder(nn.Module):
    def __init__(self, en_vocab_size, emb_dim, hid_dim, n_layers, dropout):
        super().__init__()
        self.hid_dim = hid_dim
        self.n_layers = n_layers
        self.embedding = nn.Embedding(en_vocab_size, emb_dim)
        self.rnn = nn.GRU(emb_dim, hid_dim, n_layers, dropout = dropout, batch_first=True, bidirectional=True)
        self.dropout = nn.Dropout(dropout)

    def forward(self, input):
        # input = [batch size, sequence len, en_vocab_size]
        embedding = self.embedding(input)
        outputs, hidden = self.rnn(self.dropout(embedding))
        # outputs = [batch size, sequence len, hid_dim * directions]
        # hidden =  [n_layers * directions, batch size  , hid_dim]

        return outputs, hidden
    
class Decoder(nn.Module):
    def __init__(self, cn_vocab_size, emb_dim, hid_dim, n_layers, dropout):
        super().__init__()
        self.cn_vocab_size = cn_vocab_size
        self.hid_dim = hid_dim * 2 #因为 Encoder 是双向的
        self.n_layers = n_layers
        self.embedding = nn.Embedding(cn_vocab_size, emb_dim)
        self.input_dim = emb_dim
        self.rnn = nn.GRU(self.input_dim, self.hid_dim, n_layers, dropout = dropout, batch_first=True)
        self.embedding2vocab1 = nn.Linear(self.hid_dim, self.hid_dim * 2)
        self.embedding2vocab2 = nn.Linear(self.hid_dim * 2, self.hid_dim * 4)
        self.embedding2vocab3 = nn.Linear(self.hid_dim * 4, self.cn_vocab_size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, input, hidden):
        # input = [batch size, cn_vocab_size]
        # hidden = [batch size, n_layers * directions, hid_dim]
        # Decoder 是单向，所以 directions=1
        input = input.unsqueeze(1)
        embedded = self.dropout(self.embedding(input))
        # embedded = [batch size, 1, emb_dim]
        output, hidden = self.rnn(embedded, hidden)
        # output = [batch size, 1, hid_dim]
        # hidden = [n_layers, batch size, hid_dim]

        # 将 RNN 的输出向量的维数转换到target语言的字典大小
        output = self.embedding2vocab1(output.squeeze(1))
        output = self.embedding2vocab2(output)
        prediction = self.embedding2vocab3(output)
        # prediction = [batch size, vocab size]
        return prediction, hidden

可以看出，Encoder 是一个双向四层的GRU，Decoder 是一个单向四层的GRU。Encoder 里 RNN 的隐藏状态输出用来初始化 Decoder 里的 RNN，Decoder 的输出用来计算损失值 loss 。

定义 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, input, target, teacher_forcing_ratio):
        # input  = [batch size, input len, vocab size]
        # target = [batch size, target len, vocab size]
        # teacher_forcing_ratio 是使用正解训练的概率
        batch_size = target.shape[0]
        target_len=target.shape[1]
        vocab_size = self.decoder.cn_vocab_size
        
        # 准备一个tensor存储输出
        outputs = torch.zeros(batch_size, target_len, vocab_size).to(self.device)
        
        encoder_outputs, hidden = self.encoder(input)
        # Encoder 最后的隐藏状态(hidden state)用来初始化 Decoder
        # 因为 Encoder 是双向RNN，所以需要对同一层两个方向的 hidden state 进行拼接
        # hidden = [num_layers * directions, batch size, hid dim] --> [num_layers, directions, batch size, hid dim]
        hidden = hidden.view(self.encoder.n_layers, 2, batch_size, -1)
        # hidden = [num_layers, batch size, hid dim * 2]
        hidden = torch.cat((hidden[:, -2, :, :], hidden[:, -1, :, :]), dim=2)
        
        dec_input = target[:, 0] # 'BOS'
        preds = []
        for t in range(1, target_len):
            output, hidden = self.decoder(dec_input, hidden)
            outputs[:, t] = output
            # 决定是否用正解来训练
            teacher_force = random.random() <= teacher_forcing_ratio
            # 取出输出概率最大的词
            top1 = output.argmax(1)
            # teacher force 为 True 用正解训练，否则用预测到的最大概率的词训练
            dec_input = target[:, t] if teacher_force and t < target_len else top1
            preds.append(top1.unsqueeze(1))
        preds = torch.cat(preds, 1)
        return outputs, preds
    def inference(self, input, target):
        # 测试模型
        # input  = [batch size, input len, vocab size]
        # target = [batch size, target len, vocab size]
        batch_size = input.shape[0]
        input_len = input.shape[1]
        vocab_size = self.decoder.cn_vocab_size
        
        outputs = torch.zeros(batch_size, input_len, vocab_size).to(self.device)
        encoder_outputs, hidden = self.encoder(input)
        # Encoder 最后的隐藏状态(hidden state)用来初始化 Decoder
        # 因为 Encoder 是双向RNN，所以需要对同一层两个方向的 hidden state 进行拼接
        # hidden = [num_layers * directions, batch size, hid dim] --> [num_layers, directions, batch size, hid dim]
        hidden = hidden.view(self.encoder.n_layers, 2, batch_size, -1)
        hidden = torch.cat((hidden[:, -2, :, :], hidden[:, -1, :, :]), dim=2)
        
        dec_input = target[:, 0] # 'BOS'
        preds = []
        for t in range(1, input_len):
            output, hidden = self.decoder(dec_input, hidden)
            outputs[:, t] = output
            # 取出输出概率最大的词
            top1 = output.argmax(1)
            # 用预测到的最大概率的词训练
            dec_input = top1
            preds.append(top1.unsqueeze(1))
        preds = torch.cat(preds, 1)
        return outputs, preds
    
encoder = Encoder(en_vocab_size, emb_dim, hid_dim, n_layers, dropout)
decoder = Decoder(cn_vocab_size, emb_dim, hid_dim, n_layers, dropout)
model = Seq2Seq(encoder, decoder, device).to(device)
print(model)
loss_function = nn.CrossEntropyLoss(ignore_index=0).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr = learning_rate)
print(optimizer)
print('num of parameters: ', sum(p.numel() for p in model.parameters() if p.requires_grad))

可以看出，Decoder 是逐词预测的，训练时使用 teacher_forcing 机制，验证和测试时不使用。损失函数中有一个参数 ignore_index=0 ，表示忽略预测分类为 0 时的 loss，因为此处分类为 0 时表示预测词为 '<PAD>' 。

定义一些有用的函数，用来存储、读取网络参数，计算 bleu score，将预测结果转为文字，创建训练 batch 的迭代器

def save_model(model, optimizer, store_model_path, step):
    torch.save(model.state_dict(), '{}/model_{}.ckpt'.format(store_model_path,step))
    return

def load_model(model, load_model_path):
    print('Load model from {}'.format(load_model_path))
    model.load_state_dict(torch.load('{}.ckpt'.format(load_model_path)))
    return model

def computebleu(sentences, targets):
    score = 0 
    assert (len(sentences) == len(targets))

    def cut_token(sentence):
        tmp = []
        for token in sentence:
            if token == '<UNK>' or token.isdigit() or len(bytes(token[0], encoding='utf-8')) == 1:
                tmp.append(token)
            else:
                tmp += [word for word in token]
        return tmp 

    for sentence, target in zip(sentences, targets):
        sentence = cut_token(sentence)
        target = cut_token(target)
        score += sentence_bleu([target], sentence, weights=(1, 0, 0, 0))
    return score

def tokens2sentence(outputs, int2word):
    sentences = []
    for tokens in outputs:
        sentence = []
        for token in tokens:
            word = int2word[str(int(token))]
            if word == '<EOS>':
                break
            sentence.append(word)
        sentences.append(sentence)
    return sentences

def infinite_iter(data_loader):
    it = iter(data_loader)
    while True:
        try:
            ret = next(it)
            yield ret
        except StopIteration:
            it = iter(data_loader)

训练与验证，每300轮 batch 训练进行验证，并存储模型

model.train()
model.zero_grad()
train_losses, val_losses, val_bleu_scores = [], [], []
loss_sum = 0.0
train_iter = infinite_iter(train_loader)

for step in range(summary_steps):
    model.train()
    sources, targets = next(train_iter)
    sources, targets = sources.to(device), targets.to(device)
    outputs, preds = model(sources, targets, teacher_forcing_ratio)
    # targets 的第一个 token 是 '<BOS>' 所以忽略
    outputs = outputs[:, 1:].reshape(-1, outputs.size(2))
    targets = targets[:, 1:].reshape(-1)
    loss = loss_function(outputs, targets)
    
    optimizer.zero_grad()
    loss.backward()
    grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1)
    optimizer.step()

    loss_sum += loss.item()
    if (step + 1) % 10 == 0:
        loss_sum = loss_sum / 10
        print ("\r","train [{}] loss: {:.3f}, Perplexity: {:.3f}".format(step+1, loss_sum, np.exp(loss_sum)), end=" ")
        train_losses.append(loss_sum)
        loss_sum = 0.0
        
        if (step + 1) % 300 == 0:
            # 每300轮 batch 训练进行验证，并存储模型
            model.eval()
            loss_val, bleu_val= 0.0, 0.0
            n = 0
            for sources_val, targets_val in val_loader:
                sources_val, targets_val = sources_val.to(device), targets_val.to(device)
                batch_size = sources_val.size(0)
                #print(batch_size)
                outputs_val, preds_val = model.inference(sources_val, targets_val)
                # targets 的第一个 token 是 '<BOS>' 所以忽略
                outputs_val = outputs_val[:, 1:].reshape(-1, outputs_val.size(2))
                targets_val = targets_val[:, 1:].reshape(-1)
                loss = loss_function(outputs_val, targets_val)
                loss_val += loss.item()
                
                # 将预测结果转为文字
                targets_val = targets_val.view(sources_val.size(0), -1)
                preds_val = tokens2sentence(preds_val, int2word_cn)
                sources_val = tokens2sentence(sources_val, int2word_en)
                targets_val = tokens2sentence(targets_val, int2word_cn)
                # 计算 Bleu Score
                bleu_val += computebleu(preds_val, targets_val)
                n += batch_size
            loss_val = loss_val/len(val_loader)
            bleu_val = bleu_val/n
            val_losses.append(loss_val)
            val_bleu_scores.append(bleu_val)
            print ("\n", "val [{}] loss: {:.3f}, Perplexity: {:.3f}, bleu score: {:.3f} ".format(step+1, loss_val, np.exp(loss_val), bleu_val))

            # 储存模型
            save_model(model, optimizer, store_model_path, step+1)
        

测试并保存结果

load_model_path = "./ckpt/model_6000"      # 读取模型位置

model = load_model(model, load_model_path) # 读取模型
model.to(device)
model.eval()
# 测试模型
loss_test, bleu_test= 0.0, 0.0
n = 0
result = []
for sources_test, targets_test in test_loader:
    sources_test, targets_test = sources_test.to(device), targets_test.to(device)
    batch_size = sources_test.size(0)
    # print(batch_size)
    outputs_test, preds_test = model.inference(sources_test, targets_test)
    # targets 的第一个 token 是 '<BOS>' 所以忽略
    outputs_test = outputs_test[:, 1:].reshape(-1, outputs_test.size(2))
    targets_test = targets_test[:, 1:].reshape(-1)
    loss = loss_function(outputs_test, targets_test)
    loss_test += loss.item()
                
    # 将预测结果转为文字
    targets_test = targets_test.view(sources_test.size(0), -1)
    preds_test = tokens2sentence(preds_test, int2word_cn)
    sources_test = tokens2sentence(sources_test, int2word_en)
    targets_test = tokens2sentence(targets_test, int2word_cn)
    for source, pred, target in zip(sources_test, preds_test, targets_test):
        result.append((source, pred, target))
    # 计算 Bleu Score
    bleu_test += computebleu(preds_test, targets_test)
    n += batch_size
loss_test = loss_test/len(test_loader)
bleu_test = bleu_test/n
print ('test loss: {}, bleu_score: {}'.format(loss_test,bleu_test))
# 储存结果
with open('./test_output.txt', 'w') as f:
    for line in result:
        print (line, file=f)