引入Attention的Seq2Seq模型-机器翻译

在上一篇博客中，我们介绍了 基于Seq2Seq模型的机器翻译。今天，在此基础上，对模型加入注意力机制Attention。

模型结构

首先，我们先了解一下模型的结构，基本与Seq2Seq模型大体一致：

1. 首先，第一部分是编码器Encoder，它接收source sentence，然后Encoder将其转换为一个包含具体语义信息的中间向量；
2. 加入Attention机制，接收Encoder的输出和隐藏层，得到attention weights、context vector和attention vector；
3. 最后，还是是一个解码器Decoder，接收Attention的输出，最后通过Decoder输出最终的翻译结果。

具体思路

第一步：将source input words即待英语单词，进行词向量Embedding；

第二步：词向量作为Encoder的输入，Encoder通常是一个RNN、LSTM或GRU的网络；

第三步：Encoder计算隐藏层state和所有输出
（这里注意不是最后一个输出，而是全部输出，因为我们要计算每个输出的attention权重，每个输出其实可以理解为对应一个单词token），作为Attention的输入，计算输出attention weights、context vector；

第四步：target input words即法语单词，同样进行进行Embedding，同样需要做预处理：在句子的最前面加上开始标识，句子的最后加上结束标识，如 “Je suis étudiant” ===>>>“<s> Je suis étudiant </s>”；

第四步：将target input words的词向量与context vector进行拼接，作为Decoder的输入；

第五步：同样"teacher forcing"的方法，将target input words向后偏移一个timestamp，
即原来
“<s> Je suis étudiant </s>”
变为
“Je suis étudiant </s>”，
得到target output words，这便是模型训练样本的真实标签；

最后：但是这里做法不同，Decoder每次只输出一个单词token的概率分布，所以需要与target output words逐个计算loss累加。

Attention的公式如下：

模型推理

引入注意力机制之后，在模型推理就不再使用teacher forcing了，因为模型最终每次只预测一个单词token，所以只需要一直进行predict直到预测单词为结束标记即可。

所以，

1. 首先，还是将待翻译文本source input words的词向量作为Encoder的输入，然后计算得到所有输出和隐藏层state；
2. Attention根据Encoder的所有输出和state计算得到context vector；
3. 这时我们会将开始标识即"<s>"的词向量与context vector拼接作为Decoder的输入；
4. 然后Decoder预测下一个单词是什么，并输出隐藏层state；
5. 继续利用原先的Encoder所有输出和本次Decoder的state传给Attention计算得到新的context vector;
6. 接着，预测的单词的词向量与新的context vector拼接又作为下一轮Decoder的输入；
7. 如此循环，直到预测的单词为结束标识即"</s>"。
8. 最后，将所有预测的单词拼接起来即可。

代码实现

代码基本与 基于Seq2Seq模型的机器翻译一致，这里只列出不同的网络结构代码。

import numpy as np
import tensorflow as tf
from sklearn.model_selection import train_test_split
from google.colab import drive
import os


class Encoder(tf.keras.Model):
    def __init__(self, vocab_size, embedding_size, hidden_size):
        super(Encoder, self).__init__()
        self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_size)
        self.lstm = tf.keras.layers.LSTM(hidden_size, return_state=True, return_sequences=True)

    def call(self, inputs, **kwargs):
        x = self.embedding(inputs)
        # TODO 这里加入sequence_length会更好，用mask的形式
        output, state_h, state_c = self.lstm(x, mask=None)

        return output, state_h, state_c


class Attention(tf.keras.Model):
    def __init__(self, hidden_size):
        super(Attention, self).__init__()
        self.W1 = tf.keras.layers.Dense(hidden_size)
        self.W2 = tf.keras.layers.Dense(hidden_size)
        self.V = tf.keras.layers.Dense(1)

    def call(self, query, values):
        # query -> [batch_size, 1, hidden_size]
        query = tf.expand_dims(query, axis=1)

        # values: [batch_size, max_seq_len, hidden_size]
        score = self.V(tf.nn.tanh(
            self.W1(query) + self.W2(values)
        ))
        # max_seq_len表示输入词的长度，这样是为每一个输入分配一个权重
        attention_weights = tf.argmax(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_size, hidden_size):
        super(Decoder, self).__init__()
        self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_size)
        self.lstm = tf.keras.layers.LSTM(hidden_size, return_state=True, return_sequences=True)

        self.fc = tf.keras.layers.Dense(vocab_size, activation="softmax")

        self.attention = Attention(hidden_size)

    def call(self, inputs, **kwargs):
        enc_state = kwargs.get('encoder_state')
        enc_output = kwargs.get('encoder_output')
        context_vector, attention_weights = self.attention(enc_state, enc_output)

        # x: [batch_size, 1, hidden_size + embedding_size]
        x = self.embedding(inputs)
        x = tf.concat([context_vector, x], axis=-1)
        x = tf.expand_dims(x, axis=1)

        output, _, state_c = self.lstm(x)
        output = tf.reshape(output, shape=[-1, output.shape[2]])

        # 最终输出的shape：[batch_size, vocab_size]
        output = self.fc(output)

        return output, state_c, attention_weights