### 获取 Keras 中 MultiHeadAttention 层的注意力权重 在 `tf.keras.layers.MultiHeadAttention` 类中，通过调用 `_compute_attention()` 方法不仅能够计算出注意力输出，还能获得注意力分数（即权重）。为了提取这些注意力权重，在构建模型时可以利用 Python 的 Lambda 函数来捕获并返回这些中间变量。 下面是一个具体的例子展示如何修改标准的 Transformer 编码器以暴露其内部使用的注意力机制： ```python import tensorflow as tf from tensorflow import keras layers = keras.layers class EncoderLayer(keras.Model): def __init__(self, d_model, num_heads,dff, rate=0.1): super().__init__() self.mha = layers.MultiHeadAttention(num_heads=num_heads,key_dim=d_model) self.ffn = keras.Sequential([ layers.Dense(dff, activation='relu'), # (batch_size, seq_len, dff) layers.Dense(d_model), # (batch_size, seq_len, d_model) ]) self.layernorm1 = layers.LayerNormalization(epsilon=1e-6) self.layernorm2 = layers.LayerNormalization(epsilon=1e-6) self.dropout1 = layers.Dropout(rate) self.dropout2 = layers.Dropout(rate) def call(self, x, training=False, mask=None): attn_output, _ = self.mha( query=x, value=x, key=x, return_attention_scores=True, attention_mask=mask ) out1 = self.layernorm1(x + self.dropout1(attn_output)) ffn_output = self.ffn(out1) final_output = self.layernorm2(out1 + self.dropout2(ffn_output)) return final_output, _ ``` 在这个自定义编码器层里，当调用 `mha` 实例时设置了参数 `return_attention_scores=True` 来指示该函数除了正常的输出外还要提供注意力得分。这使得可以在每一层之后访问到对应的注意力权重[^1]。 对于想要单独获取某一层或多层的注意力权重的情况，则可以通过调整上述代码结构或者创建额外的方法来进行处理。例如，如果只关心最后一层的结果，可以直接保存最后一个 `attn_output` 变量；而如果是希望得到所有层次的信息，则可能需要维护一个列表或者其他数据结构来存储每一步产生的注意力权重[^2]。