Keras:如何在LSTM模型中显示注意力权重

import tensorflow as tf
from tensorflow.keras.layers import Layer, LSTM, Dense, Input
from tensorflow.keras.models import Model

class Attention(Layer):
    def __init__(self, **kwargs):
        super(Attention, self).__init__(**kwargs)

    def build(self, input_shape):
        self.W = self.add_weight(name="att_weight", shape=(input_shape[-1], 1), initializer="normal")
        self.b = self.add_weight(name="att_bias", shape=(input_shape[1], 1), initializer="zeros")
        super(Attention, self).build(input_shape)

    def call(self, x):
        e = tf.matmul(x, self.W) + self.b
        a = tf.nn.softmax(e, axis=1)
        output = x * a
        return tf.reduce_sum(output, axis=1)

# 输入维度
input_dim = 10
# LSTM单元数
lstm_units = 64
# 输出维度
output_dim = 1

# 输入层
inputs = Input(shape=(None, input_dim))
# LSTM层
lstm_out = LSTM(lstm_units, return_sequences=True)(inputs)
# 注意力层
attention_output = Attention()(lstm_out)
# 输出层
outputs = Dense(output_dim, activation='sigmoid')(attention_output)

# 构建模型
model = Model(inputs=inputs, outputs=outputs)

# 编译模型
model.compile(optimizer='adam', loss='binary_crossentropy')

# 打印模型结构
model.summary()