当前位置: 首页>>代码示例 >>用法及示例精选 >>正文


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


MultiHeadAttention 层。

继承自:LayerModule

用法

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
)

参数

  • num_heads 注意头的数量。
  • key_dim 查询和键的每个注意力头的大小。
  • value_dim 每个注意力头的价值大小。
  • dropout 辍学概率。
  • use_bias 布尔值,密集层是否使用偏置向量/矩阵。
  • output_shape 输出张量的预期形状,除了批次和序列暗淡。如果未指定,则投影回关键函数暗淡。
  • attention_axes 应用注意力的轴。 None 表示对所有轴的注意力,但批处理、头部和特征。
  • kernel_initializer 密集层内核的初始化程序。
  • bias_initializer 密集层偏差的初始化程序。
  • kernel_regularizer 密集层内核的正则化器。
  • bias_regularizer 密集层偏差的正则化器。
  • activity_regularizer 密集层活动的正则化器。
  • kernel_constraint 密集层内核的约束。
  • bias_constraint 密集层内核的约束。

返回

  • attention_output 计算结果,形状为 (B, T, E) ,其中 T 用于目标序列形状,如果 output_shapeNone ,则 E 是查询输入的最后一个维度。否则,multi-head 输出将投影到 output_shape 指定的形状。
  • attention_scores [可选] multi-head 注意力轴上的注意力系数。

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

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

相关用法


注:本文由纯净天空筛选整理自tensorflow.org大神的英文原创作品 tf.keras.layers.MultiHeadAttention。非经特殊声明,原始代码版权归原作者所有,本译文未经允许或授权,请勿转载或复制。