从零实现基于sequence2sequence的中英翻译模型（torch版）

前言

初学torch，复现了一波官网的tutorial的聊天机器人，只不过把任务场景换成了中英翻译并且简化了一些步骤，力求做到对初学者友好，如果是初学nlp，那这个案例将是一个很好的入门案例。官网链接在此。

这篇博客仅对代码做一个记录和简单的说明，不涉及算法的数学原理。阅读此博客需要的知识储备有：
1、nlp中的基本概念，如word embedding
2、sequence2sequence架构的原理，可以通过原论文来学习。
3、GRU原理

数据来源

数据来源：http://www.manythings.org/anki/
包含包含了20133条中英文翻译

数据如图所示，第一列为英文，第二列为对应的中文，第三列不知道是什么属性，反正没用。

代码总结

import

需要以下的包

from torch import nn
import torch
import torch.nn.functional as F
import pandas as pd
import numpy as np
from sklearn.model_selection import KFold
import random
import itertools
import jieba

词库建立

一般nlp任务都需要有个这，负责词到index编号的映射，以方便后续通过index再次映射到embeding，也方便softmax输出后通过index找到对应词。

USE_CUDA = torch.cuda.is_available()
device = torch.device("cuda:0" if USE_CUDA else "cpu")
PAD_token = 0
SOS_token=1
EOS_token=2

#定义词库
class Voc:
    def __init__(self, name):
        self.name = name
        self.word2index = {}
        self.index2word = {PAD_token: "PAD", SOS_token: "/t", EOS_token: "/n"}   #/t开头，/n结尾
        self.n_words = 3

    def addSentence(self, sentence):   #setence为一个句子分词后的list
        for word in sentence:
            self.addWord(word)

    def addWord(self, word):
        if word not in self.word2index:
            self.word2index[word] = self.n_words
            self.index2word[self.n_words] = word
            self.n_words += 1

    def setence2index(self,setence): #setence为一个句子分词后的list
        index=[]
        for word in setence:
            index.append(self.word2index[word])
        return index

其中SOS_token和EOS_token为起始符和休止符，PAD_token为解决一个batch中句子长短不齐的padding符。

Voc类的功能在于建立word2index和index2word两个字典，并将sentence转换成index。

搭建编/解码器

#定义编码器
class EncoderGRU(nn.Module):
    def __init__(self,hidden_size,embedding):
        super(EncoderGRU, self).__init__()
        self.hidden_size = hidden_size
        self.embedding = embedding
        self.gru = nn.GRU(hidden_size, hidden_size)

    def forward(self, input,input_len,hidden):
        embedded = self.embedding(input)
        packed = nn.utils.rnn.pack_padded_sequence(embedded, input_len)
        output, hidden = self.gru(packed, hidden)
        output, _ = nn.utils.rnn.pad_packed_sequence(output)
        return output, hidden
        
#定义解码器
class DecoderGRU(nn.Module):
    def __init__(self, hidden_size,output_size,embedding):
        super(DecoderGRU, self).__init__()
        self.hidden_size = hidden_size
        self.embedding = embedding
        self.gru = nn.GRU(hidden_size, hidden_size)
        self.out = nn.Linear(hidden_size, output_size)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, input, hidden):
        output = self.embedding(input)
        output = F.relu(output)
        output, decoder_hidden = self.gru(output, hidden)
        output = self.softmax(self.out(output[0]))
        return output, decoder_hidden

其中nn.utils.rnn.pack_padded_sequence和nn.utils.rnn.pad_packed_sequence分别是对一个含有不等长句子的batch的压缩和解压缩。

数据预处理与loss遮盖处理

一个batch中的句子长短不齐，编码器如上一节所示一样padding后进行压缩和解压即可。但解码器涉及到loss的计算，需要利用mask矩阵对无效的地方进行遮盖，仅计算有效的loss。

#padding函数，用于对不等长的句子进行填充，返回填充后的转置
def zeroPadding(l, fillvalue=PAD_token):
    return list(itertools.zip_longest(*l, fillvalue=fillvalue))

#制作mask矩阵
def binaryMatrix(l, value=PAD_token):
    m = []
    for i, seq in enumerate(l):
        m.append([])
        for token in seq:
            if token == value:
                m[i].append(0)
            else:
                m[i].append(1)
    return m

#对loss进行mask操作
def maskNLLLoss(inp, target, mask):
    # 收集目标词的概率，并取负对数
    crossEntropy = -torch.log(torch.gather(inp, 1, target.view(-1, 1)).squeeze(1))
    # 只保留mask中值为1的部分，并求均值
    loss = crossEntropy.masked_select(mask).mean()
    loss = loss.to(device)
    return loss

#将sentence转换成index
def sentence2index_eng(voc,sentence):
    return [voc.word2index[word] for word in sentence.split()] + [EOS_token]
def sentence2index_chi(voc,sentence):
    return [voc.word2index[word] for word in jieba.lcut(sentence)] + [EOS_token]

#输入处理函数
def input_preprocecing(l,voc):  #接受一个batch的输入
    index_batch=[sentence2index_eng(voc,sentence) for sentence in l]
    lengths=torch.tensor([len(index) for index in index_batch])
    padList=zeroPadding(index_batch)
    padVar=torch.LongTensor(padList)
    return padVar,lengths

#输出处理函数
def output_preprocecing(l,voc):
    index_batch=[sentence2index_chi(voc,sentence) for sentence in l]
    max_label_len=max([len(index) for index in index_batch])
    padList=zeroPadding(index_batch)
    mask=binaryMatrix(padList)
    mask=torch.BoolTensor(mask)
    padVar=torch.LongTensor(padList)
    return padVar,mask,max_label_len

#数据预处理总函数
#接受一个batch的成对儿的数据[['...','...'],[],[],[]]
def data_preprocecing(voc_e,voc_c,pair_batch):
    pair_batch.sort(key=lambda x: len(x[0].split()),reverse=True) #按长度排序
    input_batch,output_batch=[],[]
    for pair in pair_batch:
        input_batch.append(pair[0])
        output_batch.append(pair[1])
    inp,lengths=input_preprocecing(input_batch,voc_e)
    output,mask,max_label_len=output_preprocecing(output_batch,voc_c)
    return inp,lengths,output,mask,max_label_len

训练函数

def train(input_var,input_len,label_var,max_label_len,mask,encoder,decoder,encoder_optimizer,decoder_optimizer,batch_size):
    #梯度归零
    encoder_optimizer.zero_grad()
    decoder_optimizer.zero_grad()
    #set device
    input_var=input_var.to(device)
    label_var=label_var.to(device)
    mask=mask.to(device)
    #初始化一些变量
    loss=0
    encoder_hidden=torch.zeros(1, batch_size, encoder.hidden_size, device=device) #初始化编码器的隐层
    #编码器前向传播
    _,encoder_hidden=encoder(input_var,input_len,encoder_hidden)
    #解码器前向传播
    decoder_input=torch.LongTensor([[SOS_token for _ in range(batch_size)]]).to(device)
    decoder_hidden=encoder_hidden
    for i in range(max_label_len):
        decoder_output,decoder_hidden=decoder(decoder_input,decoder_hidden)
        topv,topi=decoder_output.topk(1)
        decoder_input=torch.LongTensor([[topi[j][0] for j in range(batch_size)]]).to(device)   #每个都取出了最大的,shape=[[tensor1,tensor2,tensor3...]]
        mask_loss=maskNLLLoss(decoder_output,label_var[i],mask[i])
        loss+=mask_loss
    loss.backward()
    encoder_optimizer.step()
    decoder_optimizer.step()
    return loss.item()

其中batch的实现是利用random随机抽取，但严格来说这样就无法实现epoch，官网也是这种做法。我找了其他实现方法，对于非等长的数据来说，貌似还没什么好的办法。整齐的数据，利用dataloader函数可轻松实现。

模型建立、训练、保存

读入数据并建立词库

english=Voc('英文')
chinese=Voc('中文')

df=pd.read_table('/Users/zhaoduidemac/python wordspace/nlp学习/sequence2sequence/cmn-eng/cmn.txt',header=None).iloc[:,:-1]
df.reset_index(drop=True, inplace=True)
df.columns=['inputs','targets']
df_pair=[]

max_chi_len=0

#往两个VOC中加句子,并形成pair
for i in range(len(df['inputs'])):
    eng=df['inputs'][i].split()
    chi=jieba.lcut(df['targets'][i])
    max_chi_len=max(0,len(chi))
    english.addSentence(eng)
    chinese.addSentence(chi)
    df_pair.append([df['inputs'][i],df['targets'][i]])

划分训练集和测试集，利用10折交叉运算进行一折然后break实现hold out

df_pair_array=np.array(df_pair)
kfold=KFold(n_splits=10,shuffle=False,random_state=random.seed(2020))
for train_index,test_index in kfold.split([i for i in range(len(df_pair_array))]):
    train_pair=df_pair_array[train_index].tolist()
    test_pair=df_pair_array[test_index].tolist()
    break

初始化一些参数和模块

#设置参数
hidden_size=256
batch_size=64
learning_rate=0.001
n_iteration=60000
embedding_eng=torch.nn.Embedding(len(english.index2word),hidden_size)
embedding_chi=torch.nn.Embedding(len(chinese.index2word),hidden_size)

#网络模块及损失函数
encoder=EncoderGRU(hidden_size,embedding_eng).to(device)
decoder=DecoderGRU(hidden_size,len(chinese.index2word),embedding_chi).to(device)
encoder_optimizer=torch.optim.Adam(encoder.parameters(),lr=learning_rate)
decoder_optimizer=torch.optim.Adam(decoder.parameters(),lr=learning_rate)

模型训练

for iteration in range(n_iteration):
    input_var,lengths,label_var,mask,max_label_len=data_preprocecing(english,chinese,[random.choice(train_pair) for _ in range(batch_size)])
    loss=train(input_var,lengths,label_var,max_label_len,mask,encoder,decoder,encoder_optimizer,decoder_optimizer,batch_size)
    print('Iteration:',iteration,'loss:',loss)

保存模型，利用torch最简单的保存方法，一块打包成“泡菜”

torch.save(encoder,'se2se_encoder.pt')
torch.save(decoder,'se2se_decoder.pt')

模型读取、预测

预测函数

#预测函数，这里只接收一句，不支持一个batch的预测
def predict(input_var,encoder,decoder,max_label_len,voc_chi):
    res=''
    encoder_hidden = torch.zeros(1, 1, encoder.hidden_size, device=device)  # 初始化编码器的隐层
    input_len=torch.tensor([len(input_var)])
    _, encoder_hidden = encoder(input_var, input_len, encoder_hidden)
    decoder_input=torch.LongTensor([[SOS_token]]).to(device)
    decoder_hidden=encoder_hidden
    for i in range(max_label_len):
        decoder_output, decoder_hidden = decoder(decoder_input, decoder_hidden)
        topv, topi = decoder_output.topk(1)
        decoder_input=torch.LongTensor([[topi.item()]])
        word=voc_chi.index2word[topi.item()]
        if word=='/n':
            break
        res+=word
    return res

模型的load与预测

encoder_load=torch.load('se2se_encoder.pt')
decoder_load=torch.load('se2se_decoder.pt')

for sentence in test_pair:
    input_var=torch.LongTensor([sentence2index_eng(english,sentence[0])]).view(-1,1)
    output=predict(input_var,encoder_load,decoder_load,max_chi_len,chinese)
    print('原始句子：',sentence,'预测翻译：',output)

最终结果

训练集翻译的可以（那是当然）。
测试集基本狗屁不通。据说翻译任务是nlp中对模型、数据综合质量要求最高的任务，可能还是数据太少吧。

但这套代码相对完整地实现了一个简单的nlp任务，仍然值得学习。