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(推理)。