详细的api说明可以参考我的另一篇文章 tf.keras.layers.Attention.

tf.keras.layers.Attention实现的是点乘注意力. 调用方式为:

import tensorflow as tf
import tensorflow.keras as keras
import tensorflow.keras.layers as layers

attention = layers.Attention(use_scale=False, dropout=0)
# [batch_sz, key_dims, dims] = [1, 4, 2]
enc_outputs = tf.constant([[1, 1], [2, 2], [3, 3], [4, 4]], dtype=tf.float32)  
# [batch_sz, value_dims, dims] = [1, 4]
value_mask = tf.constant([[True, True, True, True]], dtype=tf.bool)
# [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])
atten
>> <tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[3.8448246, 3.8448246]], dtype=float32)>

这里attention([dec_outputs, enc_outputs, enc_outputs], [None, value_mask])包含两组参数:

• [dec_outputs, enc_outputs, enc_outputs]: 分别是注意力的 [query, key, value], 其中 value 可以省略, 默认key=value, 即输入 [query, key].
• [None, value_mask], 第一个参数为query mask, 第二个参数为value mask.

接下来自己计算一下是否和api调用结果相同:

# 注意力分数
score = tf.reduce_sum(enc_outputs * dec_outputs, 2)
score
>> <tf.Tensor: shape=(1, 4), dtype=float32, numpy=array([[2., 4., 6., 8.]], dtype=float32)>

# 注意力权值
weight = keras.activations.softmax(score, axis=1)
weight
>> <tf.Tensor: shape=(1, 4), dtype=float32, numpy=array([[0.00214401, 0.0158422 , 0.11705891, 0.8649548 ]], dtype=float32)>

# 加权注意力
att = tf.reduce_sum(tf.expand_dims(weight, 2) * enc_outputs, 1)
att
>> <tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[3.8448246, 3.8448246]], dtype=float32)>

可以看到结果和调用api是一样的.

value maks

enc_outputs = tf.constant([[1, 1], [2, 2], [3, 3], [4, 4]], dtype=tf.float32)  # [1, 4, 2]
value_mask = tf.constant([[True, True, False, False]], dtype=tf.bool)
dec_outputs = tf.constant([[[1, 1]]], dtype=tf.float32)  # [1, 1, 2]

atten = attention([dec_outputs, enc_outputs, enc_outputs], [None, value_mask])
atten
>> <tf.Tensor: shape=(1, 1, 2), dtype=float32, numpy=array([[[1.880797, 1.880797]]], dtype=float32)>

这里加上了对value最后两个step的mask, value_mask = tf.constant([[True, True, False, False]], dtype=tf.bool)

score = tf.reduce_sum(enc_outputs* dec_outputs, 2) - 1.e9 * (1 - tf.cast(value_mask, dtype=tf.float32))
score
>> <tf.Tensor: shape=(1, 4), dtype=float32, numpy=array([[ 2.e+00,  4.e+00, -1.e+09, -1.e+09]], dtype=float32)>

weight = keras.activations.softmax(score, axis=1)
weight
>> <tf.Tensor: shape=(1, 4), dtype=float32, numpy=array([[0.11920292, 0.880797  , 0.        , 0.        ]], dtype=float32)>

att = tf.expand_dims(weight, 2) * enc_outputs
att = tf.reduce_sum(att, 1)
att
>> <tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[1.880797, 1.880797]], dtype=float32)>

可以看到结果也是和调用api一样的. 值得注意的是, tensorflow 对mask位置的处理方式是, 对需要mask的位置在计算softmax之前-1.e9.