Python深度学习基于Tensorflow（9）注意力机制

注意力机制的主要思想是将注意力集中在信息的重要部分，对重要部分投入更多的资源，以获取更多所关注目标的细节信息，抑制其他无用信息；

在注意力机制的背景下，我们将自主性提示称为查询(Query)。对于给定任何查询，注意力机制通过集中注意力(Attention Pooling)选择感官输入(Sensory Input)，这些感官输入被称为值(Value)。每个值都与其对应的非自主提示的一个键(Key)成对。通过集中注意力，为给定的查询(自主性提示)与键(非自主性提示)进行交互，从而引导选择偏向值；

注意力机制计算过程如下：

注意力机制是怎么工作的

图中有四个东西，Query，Key，Value，Attention，图中把Key和Value放在了一起是因为其产生的output是与Key，Value的维度是无关的，只与Query的维度有关；文中把Key，Value称作为Context sequence，Query还是称作为Query sequence；为什么要这么做？可以看模型中右上方的Multi-Head Attention和左下角的Multi-Head Attention的区别进行分析；

Query，Key，Value这三个东西可以用python中的字典来解释，Key，Value表示字典中的键值对，而Query表示我们需要查询的键，Query与Key，Value匹配其得到的结果就是我们需要的信息；但是在这里并不要求Query与Key严格匹配，只需要模糊匹配就可以；Query对每一个Key进行一次模糊的匹配，并给出匹配程度，越适配权重就越大，然后根据权重再与每一个Value进行组合，得到最后的结果；其匹配程度的权重就代表了注意力机制的权重；

多头注意力机制就是把Query，Key，Value多个维度的向量分为几个少数维度的向量组合，再在Query_i，Key_i，Value_i上进行操作，最后把结果合并；

注意力机制的类型

Transformer模型中Multi-Head Attention有三个， 这三个分别对应三种Multi-Head Attention Layer：the cross attention layer，the global self attention layer， the causal self attention layer，从图中也可以发现每一层都有各自的不同，下面来一一介绍；

the cross attention layer：模型右上角(解码器)的Multi-Head Attention是注意力机制最直接的使用，其将编码器的信息和解码器的信息充分的结合了起来，文中把这一层叫做the cross attention layer；其context sequence是Encoder中得到的；

在这里要注意的是，每一次的查询是可以看得到所有的Key，Value的，但是查询与查询相互之间是看不到的，即独立的；

the global self attention layer：模型左下角(编码器)的Multi-Head Attention，这一层负责处理上下文序列，并沿着他的长度去传播信息即Query与Query之间的信息；

Query与Query之间的信息传播有很多种方式，例如在Transformer没出来之间我们普遍采用Bidirectional RNNs 和 CNNs的方式来处理；

但是为什么这里不使用RNN和CNN的方法呢？

RNN和CNN的限制：

• RNN 允许信息沿着序列一路流动，但是它要经过许多处理步骤才能到达那里(限制梯度流动)。这些 RNN 步骤必须按顺序运行，因此 RNN 不太能够利用现代并行设备的优势。
• 在 CNN 中，每个位置都可以并行处理，但它只提供有限的接收场。接收场只随着 CNN 层数的增加而线性增长，需要叠加许多卷积层来跨序列传输信息(小波网通过使用扩张卷积来减少这个问题)。

the global self attention layer允许每个序列元素直接访问每个其他序列元素，只需少量操作，并且所有输出都可以并行计算。 就像下图这样：

虽然图像类似于线性层，其本质好像也是线性层，但是其信息传播能力要比普通的线性层要强；

the causal self attention layer：因果自注意层，这一层与the global self attention layer类似，其不同的特点是需要mask，Masked Multi-Head Attention；这里要注意的是Transformer是一个自回归模型，每次产生一个输出并且把输出当作输入用来继续产生输出，因此这些模型确保每个序列元素的输出仅依赖于先前的序列元素，所以需要对Attention weights进行处理；

根据论文公式:$$Attention(Q,K,V)=softmax(\frac{Q{K}^T}{\sqrt{d_k}})V$$
可以发现$Q$和$K$是做矩阵运算的，$Attention$中的单个元素计算是$Q$中的某行和$K$中的某行做点积运算；论文中是计算$Attention$后在相应掩码位置乘一个无穷小，这里其实可以优化一下，只计算$Attention$有效的位置，可以加快速度；

这一层根据训练和推理有一点不同，如上文所说，在训练时，我们不需要每一次新产生的输出做为输入，我们直接用真值(shift)代替新产生的输出进行训练就好，这样可以在缺失一点稳定性的情况下，加快训练速度，并得到每一个位置上的损失大小；

在推理时，我们并没有真值，我们只能以每一次新产生的输出作为输入进行计算，这里有两种自回归的处理方式；一个是RNN，一个是CNN：Fast Wavenet；

以下是该层的简要表示，即每个序列元素的输出仅依赖于先前的序列元素；

构建Transformer模型

模型结构如下：

Embedding层

Embedding层分为两个部分：input_embedding 和 position_embedding

首先是input_embedding，假设 “注意力机制是Transformer的核心” 这句话分词得到 [注意力机制, 是, Transformer, 的, 核心]；用序号表示为[1,2,3,4,5]，由于是分类变量，需要转化为独热编码再使用权重矩阵映射到某一维度，然而这里可以直接用tf.kears.layers.Embedding一步实现上述两步功能；其需要一个input_dim：表示分类个数，也就是序号的最大值；output_dim：表示隐藏层维度，也就是d_model

其次是position_embedding，论文中采用的是绝对位置编码的方式：

其需要两个参数，一个是pos，这是表示位置，也就是 输入序列的长度 范围内的值，第二个是i，这表示嵌入向量的位置，也就是d_model 隐藏维度范围内的值；实现代码如下：

def absolute_position_embedding(n_len, d_model):
    positions = tf.range(n_len, dtype=tf.float32)[:,tf.newaxis]
    depths = tf.range(d_model//2, dtype=tf.float32)[tf.newaxis, :]/(d_model//2)
    angle_rads = positions / (10000**depths)
    pos_embedding = tf.concat([tf.sin(angle_rads), tf.cos(angle_rads)], axis=-1)
    return pos_embedding

如果需要画图，使用下面代码：

position_input = absolute_position_embedding(n_len=2048, d_model=512)
plt.pcolormesh(position_input.numpy().T, cmap='RdBu')
plt.ylabel('Depth')
plt.xlabel('Position')
plt.colorbar()
plt.show()

结合起来获得Embeding，使用 tf.keras.layers.Layer 类表达如下：

class Embedding(tf.keras.layers.Layer):
    def __init__(self, n_len, input_dim, d_model):
        super(Embedding, self).__init__()
        self.input_embedding = tf.keras.layers.Embedding(input_dim=input_dim, output_dim=d_model, mask_zero=True)
        self.position_embedding = self.absolute_position_embedding(n_len=n_len, d_model=d_model)
        
    # 自动添加mask
    def compute_mask(self, *args, **kwargs):
        return self.input_embedding.compute_mask(*args, **kwargs)
        
    def absolute_position_embedding(self, n_len, d_model):
        positions = tf.range(n_len, dtype=tf.float32)[:,tf.newaxis]
        depths = tf.range(d_model//2, dtype=tf.float32)[tf.newaxis, :]/(d_model//2)
        angle_rads = positions / (10000**depths)
        pos_embedding = tf.concat([tf.sin(angle_rads), tf.cos(angle_rads)], axis=-1)
        return pos_embedding

    def call(self, x):
        length = tf.shape(x)[1]
        x = tf.cast(x, tf.float32)
        x = self.input_embedding(x)
        output = x + self.position_embedding[tf.newaxis, :length]
        return output

注意力机制的实现

由于要适应三种注意力机制，我们需要分离 mask，q，k，v功能，在这里要注意，不能使用tensor.shape，而要使用tf.shape(tensor)

class MultiHeadAttention(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()

        self.num_heads = num_heads
        self.d_model = d_model
        self.wq = tf.keras.layers.Dense(d_model)
        self.wk = tf.keras.layers.Dense(d_model)
        self.wv = tf.keras.layers.Dense(d_model)

        self.wo = tf.keras.layers.Dense(d_model)

    def split_head(self, x):
        """(batch_size, heads, len, depth)"""
        output = tf.reshape(x, [tf.shape(x)[0], tf.shape(x)[1], self.num_heads, -1])
        output = tf.transpose(output, [0,2,1,3])
        return output

    def scaled_dot_product_attention(self, q, k, v, use_casual_mask):
        dk = tf.cast(tf.shape(k)[-1], tf.float32)
        matmul_qk = tf.matmul(q, k, transpose_b=True) / tf.math.sqrt(dk)

        if use_casual_mask:
            casual_mask = 1 - tf.linalg.band_part(tf.ones_like(matmul_qk), -1, 0)
            matmul_qk += casual_mask * -1e9
    
        attention_weighs = tf.nn.softmax(matmul_qk, axis=-1)
        outputs = tf.matmul(attention_weighs, v)
        return attention_weighs, outputs
        
    def call(self, q, k, v,  return_attention_weights=False, use_casual_mask=False):
        q = self.split_head(self.wq(q))
        k = self.split_head(self.wk(k))
        v = self.split_head(self.wv(v))
        attention_weighs, output = self.scaled_dot_product_attention(q, k, v, use_casual_mask=use_casual_mask)
        output = tf.transpose(output, [0,2,1,3])
        output = tf.reshape(output, [tf.shape(output)[0], tf.shape(output)[1], self.d_model])
        output = self.wo(output)
        if return_attention_weights:
            return attention_weighs, output
        else:
            return output

Encoder实现

由于 Encoder 有六个相同的子块 EncoderLayer，这里先实现一下 EncoderLayer；

EncoderLayer 中 包含一个自注意力机制以及一个前反馈层，所以需要三个参数初始化：d_model ， num_heads，dff

这里 dff 表示前反馈层中的投影维度；

class EncoderLayer(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads, dff):
        super(EncoderLayer, self).__init__()
        self.mha = MultiHeadAttention(d_model, num_heads)
        self.fnn_in = tf.keras.layers.Dense(dff, activation='relu')
        self.fnn_out = tf.keras.layers.Dense(d_model)
        self.layernorm1 = tf.keras.layers.LayerNormalization()
        self.layernorm2 = tf.keras.layers.LayerNormalization()

    def call(self, x):
        x1 = self.mha(q=x, k=x, v=x)
        x1 = self.layernorm1(x + x1)
        x2 = self.fnn_in(x1)
        x2 = self.fnn_out(x2)
        output = self.layernorm2(x1 + x2)
        return output

Encoder 从本质来说就是 EncoderLayer 的累加：

class Encoder(tf.keras.layers.Layer):
    def __init__(self, num_encoder_layer, d_model, num_heads, dff):
        super(Encoder, self).__init__()
        self.encoder_layers = [EncoderLayer(d_model, num_heads, dff) for _ in range(num_encoder_layer)]

    def call(self, x):
        for encoder_layer in self.encoder_layers:
            x = encoder_layer(x)
        return x

Decoder实现

同样按照Encoder中一样实现，但是不同的是：首先多了一个 Mask MutliHead Attention，其次第二个MHA中 k 和 v 是来自于Encoder

class Decoder_layer(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads, dff):
        super(Decoder_layer, self).__init__()
        self.mha1 = MultiHeadAttention(d_model, num_heads)
        self.mha2 = MultiHeadAttention(d_model, num_heads)

        self.fnn_in = tf.keras.layers.Dense(dff, activation='relu')
        self.fnn_out = tf.keras.layers.Dense(d_model)
        
        self.layernorm1 = tf.keras.layers.LayerNormalization()
        self.layernorm2 = tf.keras.layers.LayerNormalization()
        self.layernorm3 = tf.keras.layers.LayerNormalization()

    def call(self, x, enc_output):
        x1 = self.mha1(q=x, k=x, v=x, use_casual_mask=True)
        x1 = self.layernorm1(x + x1)
        x2 = self.mha2(q=x1, k=enc_output, v=enc_output)
        x2 = self.layernorm2(x1 + x2)
        x3 = self.fnn_in(x2)
        x3 = self.fnn_out(x3)
        output = self.layernorm3(x2 + x3)
        return output

循环 DecoderLayer 得到 Decoder：

class Decoder(tf.keras.layers.Layer):
    def __init__(self, num_decoder_layer, d_model, num_heads, dff):
        super(Decoder, self).__init__()
        self.decoder_layers = [Decoder_layer(d_model, num_heads, dff) for _ in range(num_decoder_layer)]

    def call(self, x, enc_output):
        for decoder_layer in self.decoder_layers:
            x = decoder_layer(x, enc_output)
        return x

Transformer实现

结合一下前面实现的东西，很轻松可以得到：

class Transformer(tf.keras.Model):
    def __init__(self, num_encoder_layer, num_decoder_layer, input_dim_1, input_dim_2, n_len, d_model, num_heads, dff):
        super(Transformer, self).__init__()
        self.embedding_1 = Embedding(n_len, input_dim_1, d_model)
        self.embedding_2 = Embedding(n_len, input_dim_2, d_model)
        self.encoder = Encoder(num_encoder_layer, d_model, num_heads, dff)
        self.decoder = Decoder(num_decoder_layer, d_model, num_heads, dff)
        self.final_dense = tf.keras.layers.Dense(input_dim_2, activation='softmax')

    def call(self, x):
        x1, x2 = x[0], x[1]
        x1 = self.embedding_1(x1)
        x2 = self.embedding_2(x2)
        x1 = self.encoder(x1)
        x2 = self.decoder(x2, x1)
        output = self.final_dense(x2)
        return output

导入数据并分词

这里使用的数据是粤语转普通话翻译，数据可以在我的博客资源中下载；

import tensorflow as tf
import pandas as pd
import numpy as np
import jieba
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import Whitespace
from tokenizers.processors import TemplateProcessing

def get_tokenizer(max_length, data):
    tokenizer = Tokenizer(BPE(unk_token='[UNK]'))
    trainer = BpeTrainer(special_tokens=["[PAD]", "[UNK]", "[START]", "[END]"])
    tokenizer.pre_tokenizer = Whitespace()
    tokenizer.post_processor = TemplateProcessing(
        single="[START] $0 [END]",
        special_tokens=[("[START]", 2), ("[END]", 3)],
    )
    tokenizer.enable_truncation(max_length=max_length)
    tokenizer.enable_padding(length=max_length, pad_id=0)
    tokenizer.train_from_iterator(data, trainer=trainer)
    return tokenizer


data = pd.read_csv('yue-zh.csv').dropna()
data.head()

这里结合 tokenizers 进行分词处理，采用jieba先分词，然后使用BPE算法获取此表；

yue_tokenizer = get_tokenizer(32, data['yue'].astype('string').map(lambda x: list(jieba.cut(x))))
zh_tokenizer = get_tokenizer(32, data['zh'].astype('string').map(lambda x: list(jieba.cut(x))))

yue_tokenizer.get_vocab_size(), zh_tokenizer.get_vocab_size()
(30000, 30000)

data['zh'] = data['zh'].map(lambda x: zh_tokenizer.encode(x).ids)
data['yue'] = data['yue'].map(lambda x: yue_tokenizer.encode(x).ids)

data = tf.data.Dataset.from_tensor_slices(data.values.tolist())
@tf.function
def process_data(x):
    x1, x2 = x[0], x[1]
    return (x1, x2[:-1]), x2[1:]


data = data.map(process_data).batch(32, drop_remainder=True)

train_data = data.take(1000)

初始化模型

model = Transformer(1,1,zh_tokenizer.get_vocab_size(), yue_tokenizer.get_vocab_size(), 32, 256, 8, 512)

model.compile(
    loss=tf.keras.losses.sparse_categorical_crossentropy,
    optimizer='adam',
    metrics=['accuracy']
)

history = model.fit(train_data, epochs=10)

得到训练结果：

Epoch 1/10
755/755 [==============================] - 19s 23ms/step - loss: 2.2625 - accuracy: 0.7145
Epoch 2/10
755/755 [==============================] - 17s 23ms/step - loss: 1.7673 - accuracy: 0.7420
Epoch 3/10
755/755 [==============================] - 17s 23ms/step - loss: 1.5616 - accuracy: 0.7603
Epoch 4/10
755/755 [==============================] - 17s 23ms/step - loss: 1.4048 - accuracy: 0.7720
Epoch 5/10
755/755 [==============================] - 17s 23ms/step - loss: 1.2557 - accuracy: 0.7856
Epoch 6/10
755/755 [==============================] - 17s 23ms/step - loss: 1.1231 - accuracy: 0.7986
Epoch 7/10
755/755 [==============================] - 17s 22ms/step - loss: 1.0160 - accuracy: 0.8115
Epoch 8/10
755/755 [==============================] - 17s 23ms/step - loss: 0.9229 - accuracy: 0.8246
Epoch 9/10
755/755 [==============================] - 17s 23ms/step - loss: 0.8327 - accuracy: 0.8379
Epoch 10/10
755/755 [==============================] - 17s 23ms/step - loss: 0.7516 - accuracy: 0.8504

推理

好像推理这里有一点问题，词表维度太高导致学习不好，这里只有一个不好的输出结果；

MAX_TOKENS = 32

class Translator(tf.Module):
    def __init__(self, tokenizers, transformer):
        self.tokenizers = tokenizers
        self.transformer = transformer
    
    def __call__(self, sentence, max_length=MAX_TOKENS):
        sentence = self.tokenizers.encode(sentence).ids
        sentence = tf.constant(sentence)[tf.newaxis, :]
        
        encoder_input = sentence
        
        start = tf.constant(2, dtype='int64')[tf.newaxis]
        end = tf.constant(3, dtype='int64')[tf.newaxis]
        
        # tf.TensorArray 类似于python中的列表
        output_array = tf.TensorArray(dtype=tf.int64, size=0, dynamic_size=True)
        # 在index=0的位置写入start
        output_array = output_array.write(0, start)
        
        for i in tf.range(max_length):
            output = tf.transpose(output_array.stack())
            predictions = self.transformer([encoder_input, output], training=False) # Shape `(batch_size, seq_len, vocab_size)`
            
            # 从seq_len中的最后一个维度选择last token
            predictions = predictions[:, -1:, :]  # Shape `(batch_size, 1, vocab_size)`.

            predicted_id = tf.argmax(predictions, axis=-1)
            
            # `predicted_id`加入到output_array中作为一个新的输入
            output_array = output_array.write(i+1, predicted_id[0])

            # 如果输出end就表明停止
            if predicted_id == end:
                break

        output = tf.transpose(output_array.stack())
        return zh_tokenizer.decode(output.numpy().tolist()[0])

测试：

translator = Translator(yue_tokenizer, model)
translator('嗰次喺跑馬地協和里見到有雞蛋仔格仔餅賣')

# '那次 在 跑马地 协 和里 看到 这些 友 的 地方 。'