tensorflow2: attention机制实现

代码参考网络上资料，如有侵权，可联系删除

1. 为什么进行attention

物理意义：将 Q, K 投影在不同的空间上，然后学习相似度。v 是key的内容表示。初始化 WQ,WK不同是self-attention具有泛化能力的原因，从而学习到序列中依赖语义关系。
如果不设置 W_Q、W_k、W_v，我们计算的权重很大程度上依赖于我们如何确定原始输入向量。

2. 官网实现

tf.keras.layers.Attention(
use_scale=False, **kwargs
)
输入为形状[batch_size，Tq，dim]的查询张量，形状[batch_size，Tv，dim]的值张量和形状[batch_size，Tv，dim]的键张量

计算遵循以下步骤：
计算形状为[batch_size，Tq，Tv]的分数作为查询键点积：
scores = tf.matmul（query，key，transpose_b = True）。

使用分数来计算形状为[batch_size，Tq，Tv]的分布：
distribution = tf.nn.softmax（scores）。

使用distribution创建具有形状[batch_size，Tq，dim]的值的线性组合：
返回tf.matmul（distribution，value）。

参数：
use_scale：如果为True，将创建一个标量变量来缩放scores
causal：设置为True可使解码器self-attention。添加一个罩，使位置i无法参与位置j> i。这样可以防止信息流从未来传递到过去。
dropout：attention scores下降的百分比

3. 自定义self-attention

# 缩放点注意力机制
def scaled_dot_product_attention(q, k, v, mask):"""Args:q: shape == (...., seq_len_q, depth)k: shape == (...., seq_len_k, depth)v: shape == (...., seq_len_v, depth_v)seq_len_v = seq_len_kmask : shape == (seq_len_q, seq_len_k)Returns:-  output: weighted sum-  attention_weights: weights of attention"""print('---q', tf.shape(q))print('---k', tf.shape(k))print('---v', tf.shape(v))# matmul_qk.shapes: [....., seq_len_q, seq_len_k]matmul_qk = tf.matmul(q, k, transpose_b=True)print('---qk', tf.shape(matmul_qk))dk = tf.cast(tf.shape(k)[-1], tf.float32)# key的向量维度。 为什么选择 key， 因为 向量初始化时是0，1的正太分布。 q，k 矩阵相乘时，在key的纬度上进行累加，方差变为k。因此做样的处理。# matmul_qk.shapes: [....., seq_len_q, seq_len_k]scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)if mask is not None:scaled_attention_logits += (mask * -1e9)  # 给mask 乘一个最小的值。 该值加到对应的padding=0位置，这时 softmax结果是趋近0。# attention_weights.shape == [...., seq_len_q, seq_len_k]attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # -1 表示在最后一个纬度熵进行softmaxprint('---attention_weights', tf.shape(attention_weights))# out_puts.shape == [...., seq_len_q, depth_v]out_puts = tf.matmul(attention_weights, v)print('---out_puts', tf.shape(out_puts))return out_puts, attention_weights# test dot-attention
# def print_scaled_dot_product_attention(q, k, v):
#     temp_out, temp_att = scaled_dot_product_attention(q, k, v, None)
#     return temp_out, temp_att
#
# temp_k = tf.constant([[1, 3, 0],
#                       [2, 1, 0],
#                       [3, 2, 0],
#                       [4, 1, 0]
#                       ], dtype=tf.float32)
#
# temp_v = tf.constant([[1, 2],
#                       [2, 1],
#                       [3, 2],
#                       [4, 1]
#                       ], dtype=tf.float32)   # seq_len_v = seq_len_k
#
# temp_q = tf.constant([[1, 3, 0],
#                       [2, 3, 0],
#                       [4, 1, 0]
#                       ], dtype=tf.float32)  # q 的depth 与 k的depth 需要相同
#
# temp_out, temp_att = print_scaled_dot_product_attention(temp_q, temp_k, temp_v)
#
# print(temp_out)
# print(temp_att)
## 带mask 的attention
def attention_mask():attention = layers.Attention(use_scale=False, dropout=0)## key: [batch_sz, key_dims, dims] = [1, 4, 2]enc_outputs = tf.constant([[1, 1], [2, 2], [3, 3], [4, 4]], dtype=tf.float32)# value: [batch_sz, value_dims, dims] = [1, 4, 2 ]# mask#如果给定, mask==False的位置输出为0#如果给定, mask==False的位置不会对输出产生贡献.value_mask = tf.constant([[True, True, False, False]], dtype=tf.bool)  ## query: [batch_sz, query_dims, dims] = [1, 1, 2]dec_outputs = tf.constant([[[1, 1]]], dtype=tf.float32)atten = attention([dec_outputs, enc_outputs, enc_outputs], [None, value_mask])  # query key value# 自定义验证# tf.reduce_sum(enc_outputs * dec_outputs, 2)  等价 matmul_qk = tf.matmul(q, k, transpose_b=True)score = tf.reduce_sum(enc_outputs * dec_outputs, 2) - 1.e9 * (1 - tf.cast(value_mask, dtype=tf.float32))weight = keras.activations.softmax(score, axis=1)att = tf.expand_dims(weight, 2) * enc_outputsatt = tf.reduce_sum(att, 1)# 验证 官网mask 方法与自定义mask方法结果一致#点击缩放后，发现与官网给的use_scale有差异# dk = tf.cast(tf.shape(enc_outputs)[-1], tf.float32)# score = tf.reduce_sum(enc_outputs * dec_outputs, 2)/tf.math.sqrt(dk) - 1.e9 * (1 - tf.cast(value_mask, dtype=tf.float32))# 单头self-attention
class OneHeadAttention(keras.layers.Layer):"""理论上x --Wq0 -->q0x --Wk0 --k0x --Wv0 --v0self attention， qkv 一样的？实战中：， qkv 不一样x --Wq0 -->q0x --Wk0 --k0x --Wv0 --v0小矩阵乘法，变为大矩阵的乘法。q -->Wq（大的矩阵）->Q -- split -> q0, q1, q2 [batch_size * seq_len,  depth]k， v 一样"""def __init__(self, d_model):super(OneHeadAttention, self).__init__()self.d_model = d_model  #self.WQ = keras.layers.Dense(self.d_model)self.WK = keras.layers.Dense(self.d_model)self.WV = keras.layers.Dense(self.d_model)self.dense = keras.layers.Dense(self.d_model)def call(self, q, k, v, mask):batch_size = tf.shape(q)[0]# 生成大的 q, k, v矩阵q = self.WQ(q)  # q.shape:(batch_size, seq_len_q, depth)k = self.WK(k)  # k.shape:(batch_size, seq_len_k, depth)v = self.WV(v)  # k.shape:(batch_size, seq_len_v, depth)# scaled_attention_out_puts.shape: [batch_size,  seq_len_q, depth ]# attention_weights.shape: [batch_size,  seq_len_q, seq_len_k],scaled_attention_out_puts, attention_weights = \scaled_dot_product_attention(q, k, v, mask)output = self.dense(scaled_attention_out_puts)return output, attention_weights#  test OneHeadAttention
temp_mha = OneHeadAttention(d_model=128)   # 类初始化
y = tf.random.uniform((1, 60, 64))  # [batch_size, seq_len_q, dim]  因此，输入的dim可以忽略。 最终输出的是v的dim
output, attention_weights = temp_mha(y, y, y, mask=None)  # 初始化 q, k, v ，默认调用call 函数。# MultiHeadAttention 实现
from tensorflow import keras
class MultiHeadAttention(keras.layers.Layer):"""理论上x --Wq0 -->q0x --Wk0 --k0x --Wv0 --v0self attention， qkv 一样的？实战中：， qkv 不一样x --Wq0 -->q0x --Wk0 --k0x --Wv0 --v0小矩阵乘法，变为大矩阵的乘法。q -->Wq（大的矩阵）->Q -- split -> q0, q1, q2 [batch_size * seq_len,  depth]k， v 一样"""def __init__(self, d_model, num_heads):super(MultiHeadAttention, self).__init__()self.num_heads = num_heads  # head 的个数self.d_model = d_model  # 多头最终concat 的dimassert self.d_model % self.num_heads == 0  # model的维度必须是head的整数倍self.depth = self.d_model // self.num_headsself.WQ = keras.layers.Dense(self.d_model)self.WK = keras.layers.Dense(self.d_model)self.WV = keras.layers.Dense(self.d_model)self.dense = keras.layers.Dense(self.d_model)def split_heads(self, x, batch_size):x = tf.reshape(x,(batch_size, -1, self.num_heads, self.depth))return tf.transpose(x, perm=[0, 2, 1, 3])def call(self, q, k, v, mask):batch_size = tf.shape(q)[0]# 生成大的 q, k, v矩阵q = self.WQ(q)  # q.shape:(batch_size, seq_len_q, depth)k = self.WK(k)  # k.shape:(batch_size, seq_len_k, depth)v = self.WV(v)  # k.shape:(batch_size, seq_len_v, depth)# q.shape :[batch_size, num_heads, seq_len_q, depth]q = self.split_heads(q, batch_size)# k.shape :[batch_size, num_heads, seq_len_k, depth]k = self.split_heads(k, batch_size)# v.shape :[batch_size, num_heads, seq_len_k, depth]v = self.split_heads(v, batch_size)# scaled_attention_out_puts.shape: [batch_size,  num_heads, seq_len_q, depth ]# attention_weights.shape: [batch_size, num_heads, seq_len_q, seq_len_k],scaled_attention_out_puts, attention_weights = \scaled_dot_product_attention(q, k, v, mask)# before ：scaled_attention_out_puts.shape:[batch_size, num_heads, seq_len_q,  depth]# after ：scaled_attention_out_puts.shape:[batch_size, seq_len_q, num_heads, depth]scaled_attention_out_puts = tf.transpose(scaled_attention_out_puts, perm=[0, 2, 1, 3])  ## 将num_heads 进行降纬度。 即多个头进行合并# concat_attention.shape: [batch_size, seq_len_q, d_model]concat_attention = tf.reshape(scaled_attention_out_puts, (batch_size, -1, self.d_model))output = self.dense(concat_attention)return output, attention_weights#  test MultiHeadAttention
temp_mha = MultiHeadAttention(d_model=512, num_heads=8)   # 类初始化
y = tf.random.uniform((1, 60, 256))  # [batch_size, seq_len_q, dim]  因此，输入的dim可以忽略。 最终输出的是v的dim
output, attention_weights = temp_mha(y, y, y, mask=None)  # 初始化 q, k, v ，默认调用call 函数。print(output)
print(attention_weights)

4. 通过attention api 实现

# 使用 tf.keras.layers.Attention() 进行attention
# 存在问题， 目前还没弄明白 mask 如何加， 待更新。
def attention_test():"""batch_size: batch的大小，seq_len: 序列的长度embedding_dim： 序列中单个元素的向量纬度"""# input layers# input.shape :(batch_size, seq_len)input_query_char = tf.keras.layers.Input(shape=(maxlen,), name="input_q_char")# 序列进行embedding# emb_q_char.shape: (batch_size, seq_len, embedding_dim)emb_q_char = tf.keras.layers.Embedding(input_dim=vocab_size,output_dim=embedding_dim,input_length=maxlen,name='emb_q_char')(input_query_char)print('----emb_q_char', emb_q_char.shape)# 初始化 q, k# q.shape: (batch_size, seq_len_q, embedding_dim)# k.shape: (batch_size, seq_len_k, embedding_dim)WQ = keras.layers.Dense(embedding_dim)WK = keras.layers.Dense(embedding_dim)q = WQ(emb_q_char)k = WK(emb_q_char)# 调用attention层# attention_outputs.shape :(batch_size, seq_len_q, embedding_dim)# attention_weights.shape: (batch_size, seq_len_q, seq_len_k)attention_outputs, attention_weights = tf.keras.layers.Attention()([q, k],return_attention_scores=True)print('---attention_outputs', attention_outputs.shape)print('---attention weights', attention_weights.shape)# Add&Norm :# x 与 attention 进行进行concat# 标注化处理, axis=1 ，表示在行上进行标准化处理# norm_data.shape: (batch_size, seq_len, embedding_dim)layer_normal = tf.keras.layers.LayerNormalization(epsilon=1e-6)norm_data = layer_normal(emb_q_char + attention_outputs)print('----norm_data shape', norm_data.shape)# FeedWork:前馈神经网络# diff ： 前馈神经网络的维度, 具体这个纬度如何定义需要确认下。# ff_data.shape : (batch_size, seq_len_q, embedding_dim)fcc = tf.keras.layers.Dense(1024, activation='relu', name='fcc')(norm_data)ff_data = tf.keras.layers.Dense(embedding_dim)(fcc)print("----fcc shape", fcc.shape)print("----ff data", ff_data.shape)# 结果进行 sum pool# attention_sum_pool.shape:(batch_size,embedding_dim )attention_sum_pool = tf.reduce_sum(ff_data, 1)print("----attention_sum shape", attention_sum_pool.shape)# 关于attention最后结果处理两种方法：# 方法1： 经过feed work net， 然后sum pool# 方法2： attention结果sum pool + input进行mean pool （后期再验证）# DNN 建模query_merge_vector = tf.keras.layers.concatenate([attention_sum_pool])query_vector_l1 = tf.keras.layers.Dense(100, activation='relu', name='query_d1')(query_merge_vector)query_vector_l2 = tf.keras.layers.Dense(100, activation='relu',name="query_d4", kernel_regularizer='l2')(query_vector_l1)output = tf.keras.layers.Dense(1, activation='sigmoid', name='output')(query_vector_l2)model = keras.models.Model(inputs=[input_query_char],outputs=[output])model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])model.summary()return modelif __name__ == '__main__':attention_test()