Python tf.keras.layers.MultiHeadAttention用法及代码示例

MultiHeadAttention 层。

继承自：Layer，Module

用法

tf.keras.layers.MultiHeadAttention(
    num_heads, key_dim, value_dim=None, dropout=0.0, use_bias=True,
    output_shape=None, attention_axes=None,
    kernel_initializer='glorot_uniform',
    bias_initializer='zeros', kernel_regularizer=None,
    bias_regularizer=None, activity_regularizer=None, kernel_constraint=None,
    bias_constraint=None, **kwargs
)

这是 multi-headed attention 的实现，如论文“Attention is all you Need”(Vaswani et al., 2017)中所述。如果query , key, value 相同，则为self-attention。 query 中的每个时间步都会处理 key 中的相应序列，并返回一个 fixed-width 向量。

该层首先投影 query , key 和 value 。这些(实际上)是长度为 num_attention_heads 的张量列表，其中相应的形状是 (batch_size, <query dimensions>, key_dim) , (batch_size, <key/value dimensions>, key_dim) , (batch_size, <key/value dimensions>, value_dim) 。

然后，查询和关键张量是dot-producted 并缩放。这些被softmaxed以获得注意力概率。然后通过这些概率对值张量进行插值，然后连接回单个张量。

最后，最后一维为value_dim的结果张量可以进行线性投影并返回。

在自定义层内使用 MultiHeadAttention 时，自定义层必须实现 build() 并调用 MultiHeadAttention 的 _build_from_signature() 。这样可以在加载模型时正确恢复权重。

在自定义层中使用时。

例子：

使用注意掩码在两个序列输入上执行一维cross-attention。返回额外的注意力权重。

layer = MultiHeadAttention(num_heads=2, key_dim=2)
target = tf.keras.Input(shape=[8, 16])
source = tf.keras.Input(shape=[4, 16])
output_tensor, weights = layer(target, source,
                               return_attention_scores=True)
print(output_tensor.shape)
(None, 8, 16)
print(weights.shape)
(None, 2, 8, 4)

在轴 2 和 3 上的 5D 输入张量上执行 2D self-attention。

layer = MultiHeadAttention(num_heads=2, key_dim=2, attention_axes=(2, 3))
input_tensor = tf.keras.Input(shape=[5, 3, 4, 16])
output_tensor = layer(input_tensor, input_tensor)
print(output_tensor.shape)
(None, 5, 3, 4, 16)

调用参数：

• query：询问Tensor形状的(B, T, dim).
• value：价值Tensor形状的(B, S, dim).
• key:可选键Tensor形状的(B, S, dim).如果没有给出，将使用value对彼此而言key和value，这是最常见的情况。
• attention_mask：形状的布尔掩码(B, T, S)，这会阻止对某些位置的注意。布尔掩码指定哪些查询元素可以关注哪些关键元素，1 表示关注，0 表示不关注。对于缺少的批次维度和头部维度，可能会发生广播。
• return_attention_scores:一个布尔值，指示输出是否应该是(attention_output, attention_scores)如果True， 或者attention_output如果False.默认为False.
• training:Python 布尔值，指示层应该在训练模式(添加 dropout)还是在推理模式(无 dropout)下运行。默认使用父层/模型的训练模式，如果没有父层，则默认为 False(推理)。