Keras 注意力层

适用于 Keras 的注意力层。支持 Luong 和 Bahdanau 的评分函数。

已通过 Tensorflow 2.8, 2.9, 2.10, 2.11, 2.12, 2.13 和 2.14 测试（2023 年 9 月 26 日）。

安装

PyPI

pip install attention

注意力层

Attention(
    units=128,
    score='luong',
    **kwargs
)

参数

units：整数。注意力向量 ($a_t$) 中的（输出）单元数量。
score：字符串。评分函数 $score(h_t, \bar{h_s})$。可能的值是 luong 或 bahdanau。
- Luong 的乘法风格。链接到论文。
- Bahdanau 的加法风格。链接到论文。

输入形状

形状为 (batch_size, timesteps, input_dim) 的 3D 张量。

输出形状

形状为 (batch_size, num_units) ($a_t$) 的 2D 张量。

如果你想查看注意力权重，可以参考这个例子 examples/add_two_numbers.py。

示例

import numpy as np
from tensorflow.keras import Input
from tensorflow.keras.layers import Dense, LSTM
from tensorflow.keras.models import load_model, Model

from attention import Attention


def main():
    # 虚拟数据。在此示例中没有任何需要学习的内容。
    num_samples, time_steps, input_dim, output_dim = 100, 10, 1, 1
    data_x = np.random.uniform(size=(num_samples, time_steps, input_dim))
    data_y = np.random.uniform(size=(num_samples, output_dim))

    # 定义/编译模型。
    model_input = Input(shape=(time_steps, input_dim))
    x = LSTM(64, return_sequences=True)(model_input)
    x = Attention(units=32)(x)
    x = Dense(1)(x)
    model = Model(model_input, x)
    model.compile(loss='mae', optimizer='adam')
    model.summary()

    # 训练。
    model.fit(data_x, data_y, epochs=10)

    # 测试保存/重新加载模型。
    pred1 = model.predict(data_x)
    model.save('test_model.h5')
    model_h5 = load_model('test_model.h5', custom_objects={'Attention': Attention})
    pred2 = model_h5.predict(data_x)
    np.testing.assert_almost_equal(pred1, pred2)
    print('成功。')


if __name__ == '__main__':
    main()