基于注意力的seq2seq模型

基于注意力的seq2seq模型

简介

  相较与普通的seq2seq模型，其多了一个注意力层。简单来说就是在encoder-decoder中，不再是由encoder直接将数据传给decoder，而是经过一个注意力层对encoder的输出进行加权，区分出对当前预测重要的数据，再根据这些重要的数据进行预测。

模型结构

 参考链接：https://www.tensorflow.org/tutorials/text/nmt_with_attention

 模型结构如图：

 模型说明：

  模型与文档Sequence to Sequence模型中实现的模型基本类似，在图形的最下方为模型的输入，在< s >左侧的是需要翻译的语句，右侧是翻译的结果。最上方式模型预测的结果。模型运行时左下方的需要翻译的语句会全部传入encoder模型，进行编码，右下的数据会逐个传入decoder中进行翻译，decoder在翻译的时候会先根据需要预测的词和encoder的编码计算注意力。

keras实现

  代码主要分为4部分：数据、模型、训练、预测。

数据处理

  首先是数据部分，代码如下：

import tensorflow as tf

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from sklearn.model_selection import train_test_split

import unicodedata
import re
import numpy as np
import os
import io
import time


def download():
    # 下载文件
    path_to_zip = tf.keras.utils.get_file(
        'spa-eng.zip', origin='http://storage.googleapis.com/download.tensorflow.org/data/spa-eng.zip',
        extract=True)

    path_to_file = os.path.dirname(path_to_zip) + "/spa-eng/spa.txt"
    print(path_to_file)


# 将 unicode 文件转换为 ascii
def unicode_to_ascii(s):
    return ''.join(c for c in unicodedata.normalize('NFD', s)
                   if unicodedata.category(c) != 'Mn')


def preprocess_sentence(w):
    w = unicode_to_ascii(w.lower().strip())

    # 在单词与跟在其后的标点符号之间插入一个空格
    # 例如： "he is a boy." => "he is a boy ."
    # 参考：https://stackoverflow.com/questions/3645931/python-padding-punctuation-with-white-spaces-keeping-punctuation
    w = re.sub(r"([?.!,¿])", r" \1 ", w)
    w = re.sub(r'[" "]+', " ", w)

    # 除了 (a-z, A-Z, ".", "?", "!", ",")，将所有字符替换为空格
    w = re.sub(r"[^a-zA-Z?.!,¿]+", " ", w)

    w = w.rstrip().strip()

    # 给句子加上开始和结束标记
    # 以便模型知道何时开始和结束预测
    w = '<start> ' + w + ' <end>'
    return w


# 1. 去除重音符号
# 2. 清理句子
# 3. 返回这样格式的单词对：[ENGLISH, SPANISH]
def create_dataset(path, num_examples):
    lines = io.open(path, encoding='UTF-8').read().strip().split('\n')

    word_pairs = [[preprocess_sentence(w) for w in l.split('\t')] for l in lines[:num_examples]]

    return zip(*word_pairs)


def max_length(tensor):
    return max(len(t) for t in tensor)


def tokenize(lang):
    lang_tokenizer = tf.keras.preprocessing.text.Tokenizer(filters='')
    lang_tokenizer.fit_on_texts(lang)

    tensor = lang_tokenizer.texts_to_sequences(lang)

    tensor = tf.keras.preprocessing.sequence.pad_sequences(tensor,
                                                           padding='post')

    return tensor, lang_tokenizer


def load_dataset(path, num_examples=None):
    # 创建清理过的输入输出对
    targ_lang, inp_lang = create_dataset(path, num_examples)

    input_tensor, inp_lang_tokenizer = tokenize(inp_lang)
    target_tensor, targ_lang_tokenizer = tokenize(targ_lang)

    return input_tensor, target_tensor, inp_lang_tokenizer, targ_lang_tokenizer

  首先看download方法，这个方法是用来下载翻译数据的。这里使用的是西班牙语和英语的翻译数据，下载下来的数据格式如下图所示：

  然后是数据处理。处理的入口方法是load_dataset方法，这个方法有两个参数：path和num_examples。其中path是文件的路径，num_examples是用来指定需要从文件中加载多少行数据作为模型使用的数据。默认是读取全部数据，这个文件大概有11万行数据，如果机器的内存不够大可以通过这个参数读取少量数据，来测试模型。

  在load_dataset方法中，首先是一个create_dataset方法，这个方法主要作用是从文件中读取数据；然后是两个tokenize方法，在处理文本数据的时候，一般的做法是先根据文本数据制作词表，然后通过词表将文本转换成数字，最后再将数字转换成向量。这里的tokenize的做的便是生成词表并将文本转换成数字的工作。

  在create_dataset方法中，首先就是通过io流读取数据，然后将读取到的数据按换行符分割。从上面的数据格式可知道，这里得到的就是一对相互对应的英语与西班牙语翻译。然后他又用了两个for循环遍历每行数据，并按\t拆分，将英语与西班牙语分开。同时还执行了一个preprocess_sentence方法，这个方法主要就是清理一下句子，添加一下开始符和结束符。

  在tokenize方法中，首先会创建一个Tokenizer对象，这个类是TensorFlow中用来处理文本数据的类，fit_on_texts方法会根据输入语料建立词表，然后是texts_to_sequences方法将语句转换成数字，最后是pad_sequences将长短不一的语句填充到相同的长度。

  最后通过load_dataset方法处理过的数据效果如下：

模型

  通常seq2seq模型分为编码器和解码器两个部分，这里相较于通常的模型还多了一个注意力层。其代码实现如下：

class Encoder(tf.keras.Model):
    def __init__(self, vocab_size, embedding_dim, enc_units, batch_sz):
        super(Encoder, self).__init__()
        self.batch_sz = batch_sz
        self.enc_units = enc_units
        self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
        self.gru = tf.keras.layers.GRU(self.enc_units,
                                       return_sequences=True,
                                       return_state=True,
                                       recurrent_initializer='glorot_uniform')

    def call(self, x, hidden):
        x = self.embedding(x)
        output, state = self.gru(x, initial_state=hidden)
        return output, state

    def initialize_hidden_state(self):
        return tf.zeros((self.batch_sz, self.enc_units))


class BahdanauAttention(tf.keras.layers.Layer):
    def __init__(self, units):
        super(BahdanauAttention, self).__init__()
        self.W1 = tf.keras.layers.Dense(units)
        self.W2 = tf.keras.layers.Dense(units)
        self.V = tf.keras.layers.Dense(1)

    def call(self, query, values):
        # 隐藏层的形状 == （批大小，隐藏层大小）
        # hidden_with_time_axis 的形状 == （批大小，1，隐藏层大小）
        # 这样做是为了执行加法以计算分数
        hidden_with_time_axis = tf.expand_dims(query, 1)

        # 分数的形状 == （批大小，最大长度，1）
        # 我们在最后一个轴上得到 1， 因为我们把分数应用于 self.V
        # 在应用 self.V 之前，张量的形状是（批大小，最大长度，单位）
        score = self.V(tf.nn.tanh(
            self.W1(values) + self.W2(hidden_with_time_axis)))

        # 注意力权重 （attention_weights） 的形状 == （批大小，最大长度，1）
        attention_weights = tf.nn.softmax(score, axis=1)

        # 上下文向量 （context_vector） 求和之后的形状 == （批大小，隐藏层大小）
        context_vector = attention_weights * values
        context_vector = tf.reduce_sum(context_vector, axis=1)

        return context_vector, attention_weights


class Decoder(tf.keras.Model):
    def __init__(self, vocab_size, embedding_dim, dec_units, batch_sz):
        super(Decoder, self).__init__()
        self.batch_sz = batch_sz
        self.dec_units = dec_units
        self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
        self.gru = tf.keras.layers.GRU(self.dec_units,
                                       return_sequences=True,
                                       return_state=True,
                                       recurrent_initializer='glorot_uniform')
        self.fc = tf.keras.layers.Dense(vocab_size)

        # 用于注意力
        self.attention = BahdanauAttention(self.dec_units)

    def call(self, x, hidden, enc_output):
        # 编码器输出 （enc_output） 的形状 == （批大小，最大长度，隐藏层大小）
        context_vector, attention_weights = self.attention(hidden, enc_output)

        # x 在通过嵌入层后的形状 == （批大小，1，嵌入维度）
        x = self.embedding(x)
        # x 在拼接 （concatenation） 后的形状 == （批大小，1，嵌入维度 + 隐藏层大小）
        x = tf.concat([tf.expand_dims(context_vector, 1), x], axis=-1)
        # 将合并后的向量传送到 GRU
        output, state = self.gru(x)

        # 输出的形状 == （批大小 * 1，隐藏层大小）
        output = tf.reshape(output, (-1, output.shape[2]))

        # 输出的形状 == （批大小，vocab）
        x = self.fc(output)

        return x, state, attention_weights

  这里分为了三个类：Encoder，BahdanauAttention，Decoder。其中Encoder是编码器，BahdanauAttention是注意力层，Decoder是解码器。

  这里的Encoder实现非常简单，就两层：Embedding和GRU。Embedding是词向量层，作用是将输入的数字转换成指定维度的向量。GRU是一种常用的循环神经网络。这里需要说明的是GRU的return_sequences参数和return_state参数被设置成为了True。这时模型的返回与通常情况下的返回是不同的。正常情况下rnn的返回值是最后一个单元的输出，return_sequences参数表示返回所有单元的输出。return_state表示返回存储的中间状态。在GRU中需要返回的中间状态和最后一个单元的输出是同一个值。所以这个模型的返回只有两个。

  然后是注意力层（BahdanauAttention）的实现。这里最重要的其实是传入的query和values，这两个参数在后面分析decoder和模型训练的时候会详细分析。这里从结果上来说query是GRU模型输出的中间状态，values是encoder输出的所有单元的输出。然后细看call方法内的操作：首先扩充query的维度，让他的维度和values相同，然后将其和values分别添加一个全连接层，并将两者结果相加，得到的和再来一个tanh函数，最后再来一个输出为1的全连接层。这样就得到了Encoder中每一个输出的分数，然后再用softmax将分数转换成权重。最后将权重与values相乘并求和，便得到了需要输出的context_vector。

  最后是Decoder，他相较于Encoder多了一个注意力层和全连接层。在调用时首先会根据传入的隐藏状态和Encoder的输出来计算注意力，然后将传入的代表语句的数字转换成词向量，然后将词向量和注意力连接到一起，传入到GRU中，最后在将GRU的输出展平，输出到一个全连接层，计算出下一个词的概率。

训练

def loss_function(real, pred):
    mask = tf.math.logical_not(tf.math.equal(real, 0))
    loss_ = loss_object(real, pred)

    mask = tf.cast(mask, dtype=loss_.dtype)
    loss_ *= mask

    return tf.reduce_mean(loss_)

@tf.function
def train_step(inp, targ, enc_hidden):
    loss = 0

    with tf.GradientTape() as tape:
        enc_output, enc_hidden = encoder(inp, enc_hidden)

        dec_hidden = enc_hidden

        dec_input = tf.expand_dims([targ_lang.word_index['<start>']] * BATCH_SIZE, 1)

        # 教师强制 - 将目标词作为下一个输入
        for t in range(1, targ.shape[1]):
            # 将编码器输出 （enc_output） 传送至解码器
            predictions, dec_hidden, _ = decoder(dec_input, dec_hidden, enc_output)

            loss += loss_function(targ[:, t], predictions)

            # 使用教师强制
            dec_input = tf.expand_dims(targ[:, t], 1)

    batch_loss = (loss / int(targ.shape[1]))

    variables = encoder.trainable_variables + decoder.trainable_variables

    gradients = tape.gradient(loss, variables)

    optimizer.apply_gradients(zip(gradients, variables))

    return batch_loss

path_to_file = "../data/spa-eng/spa.txt"
num_examples = 3000
input_tensor, target_tensor, inp_lang, targ_lang = load_dataset(path_to_file, num_examples)



# 计算目标张量的最大长度 （max_length）
max_length_targ, max_length_inp = max_length(target_tensor), max_length(input_tensor)

# print(max_length_inp)
# print(max_length_inp)

# 采用 80 - 20 的比例切分训练集和验证集
input_tensor_train, input_tensor_val, target_tensor_train, target_tensor_val = train_test_split(input_tensor,
                                                                                                target_tensor,
                                                                                                test_size=0.2)
BUFFER_SIZE = len(input_tensor_train)
BATCH_SIZE = 8
steps_per_epoch = len(input_tensor_train) // BATCH_SIZE
embedding_dim = 256
units = 1024
vocab_inp_size = len(inp_lang.word_index) + 1
vocab_tar_size = len(targ_lang.word_index) + 1
dataset = tf.data.Dataset.from_tensor_slices((input_tensor_train, target_tensor_train)).shuffle(BUFFER_SIZE)
dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)
encoder = Encoder(vocab_inp_size, embedding_dim, units, BATCH_SIZE)
decoder = Decoder(vocab_tar_size, embedding_dim, units, BATCH_SIZE)




optimizer = tf.keras.optimizers.Adam()
loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
    from_logits=True, reduction='none')

checkpoint_dir = '../savemodel/'
checkpoint_prefix = os.path.join(checkpoint_dir, "seq2seq_ckpt")
checkpoint = tf.train.Checkpoint(optimizer=optimizer,
                                 encoder=encoder,
                                 decoder=decoder)

EPOCHS = 10

for epoch in range(EPOCHS):
    start = time.time()

    enc_hidden = encoder.initialize_hidden_state()
    total_loss = 0

    for (batch, (inp, targ)) in enumerate(dataset.take(steps_per_epoch)):
        batch_loss = train_step(inp, targ, enc_hidden)
        total_loss += batch_loss

        if batch % 100 == 0:
            print('Epoch {} Batch {} Loss {:.4f}'.format(epoch + 1,
                                                         batch,
                                                         batch_loss.numpy()))
    # 每 2 个周期（epoch），保存（检查点）一次模型
    if (epoch + 1) % 2 == 0:
        checkpoint.save(file_prefix=checkpoint_prefix)

    print('Epoch {} Loss {:.4f}'.format(epoch + 1,
                                        total_loss / steps_per_epoch))
    print('Time taken for 1 epoch {} sec\n'.format(time.time() - start))

  这里就不详细分析加载数据、创建模型等代码了，这段代码的重点在train_step和loss_function这两个方法。

  首先是train_step方法，在训练时每次从数据集中拿出一个批次数据都会调用该方法进行训练，在调用的时候传入了三个参数：inp, targ, enc_hidden。inp和targ分别对应需要翻译的两种语言，enc_hidden是GRU的初始状态，他是一个全为零的数组。

  然后是train_step方法内部，首先他会将inp直接传入encoder中进行编码，得到两个输出enc_output和enc_hidden，然后是解码器，解码器的训练不像编码器直接将所有数据一起输入。解码器会按照单词一个一个的训练。解码器的第一个输入是固定的开始符，然后是for循环里会直接调用decoder，这里需要注意传入的三个参数：dec_input, dec_hidden, enc_output。联系上文讲解decoder时内容，这里的dec_input是call方法中的x，dec_hidden是call中的hidden，他与enc_output一起传入了注意力层中，其中dec_hidden是注意力中的query，enc_output是values。细看train_step方法可以发现传入的enc_output是一直不变的，但是dec_hidden却是会变的。他一开始是encoder的隐藏状态，然后再for中他会被解码器的隐藏状态覆盖。

  在decoder调用用结束后便会调用loss_function方法来计算损失。损失函数用的是交叉熵损失函数，但loss_function中还有一段对mask的操作。这是因为在翻译数据中每个句子的长度都是不同的，同一批次的同一个位置，有的数据可能真实的词有的可能是填充的零。mask的作用就是去除掉这些填充的零。这里的做法也很简单，首先判断该位置是否为0，然后对其取反，这样是0的位置为false，有词的位置为true。再将这些Boolean值转换成int，false为0，true为1。最后再将其与交叉熵的结果相乘。这样填充位置的损失会乘以0，变成0，有词的位置会乘以1，损失不变。

预测

def evaluate(sentence):
    attention_plot = np.zeros((max_length_targ, max_length_inp))

    sentence = preprocess_sentence(sentence)

    inputs = [inp_lang.word_index[i] for i in sentence.split(' ')]
    inputs = tf.keras.preprocessing.sequence.pad_sequences([inputs],
                                                           maxlen=max_length_inp,
                                                           padding='post')
    inputs = tf.convert_to_tensor(inputs)

    result = ''

    hidden = [tf.zeros((1, units))]
    enc_out, enc_hidden = encoder(inputs, hidden)

    dec_hidden = enc_hidden
    dec_input = tf.expand_dims([targ_lang.word_index['<start>']], 0)

    for t in range(max_length_targ):
        predictions, dec_hidden, attention_weights = decoder(dec_input,
                                                             dec_hidden,
                                                             enc_out)

        # 存储注意力权重以便后面制图
        attention_weights = tf.reshape(attention_weights, (-1,))
        attention_plot[t] = attention_weights.numpy()

        predicted_id = tf.argmax(predictions[0]).numpy()

        result += targ_lang.index_word[predicted_id] + ' '

        if targ_lang.index_word[predicted_id] == '<end>':
            return result, sentence, attention_plot

        # 预测的 ID 被输送回模型
        dec_input = tf.expand_dims([predicted_id], 0)

    return result, sentence, attention_plot


# 注意力权重制图函数
def plot_attention(attention, sentence, predicted_sentence):
    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(1, 1, 1)
    ax.matshow(attention, cmap='viridis')

    fontdict = {
    'fontsize': 14}

    ax.set_xticklabels([''] + sentence, fontdict=fontdict, rotation=90)
    ax.set_yticklabels([''] + predicted_sentence, fontdict=fontdict)

    ax.xaxis.set_major_locator(ticker.MultipleLocator(1))
    ax.yaxis.set_major_locator(ticker.MultipleLocator(1))

    plt.show()


def translate(sentence):
    result, sentence, attention_plot = evaluate(sentence)

    print('Input: %s' % (sentence))
    print('Predicted translation: {}'.format(result))

    attention_plot = attention_plot[:len(result.split(' ')), :len(sentence.split(' '))]
    plot_attention(attention_plot, sentence.split(' '), result.split(' '))


# 恢复检查点目录 （checkpoint_dir） 中最新的检查点
checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
translate(u'hace mucho frio aqui.')
translate(u'esta es mi vida.')

  预测的方法与训练的大致相同，不同的地方在于没有解码器需要将自身预测的词作为输入来预测下一个词。