直接上干货

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

# 定义Transformer模型
def transformer_model(input_shape, num_layers, d_model, num_heads, dff, dropout_rate):
    inputs = layers.Input(shape=input_shape)

    # 添加掩码
    padding_mask = layers.Lambda(lambda x: tf.cast(tf.math.equal(x, 0), tf.float32))(inputs)
    encoder_masking = layers.Masking(mask_value=0.0)(inputs)

    # 编码器
    x = layers.Dense(d_model, activation="relu")(encoder_masking)
    for i in range(num_layers):
        x = layers.MultiHeadAttention(num_heads=num_heads, key_dim=d_model)(
            [x, x],
            attention_mask=[padding_mask, None],
        )
        x = layers.Dropout(dropout_rate)(x)
        x = layers.LayerNormalization(epsilon=1e-6)(x)

        # 前馈网络
        ffn = keras.Sequential(
            [layers.Dense(dff, activation="relu"), layers.Dense(d_model)]
        )
        x = ffn(x)
        x = layers.Dropout(dropout_rate)(x)
        x = layers.LayerNormalization(epsilon=1e-6)(x)

    # 解码器
    outputs = layers.Dense(1)(x)
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model

# 定义输入数据
input_shape = (10, 1) # 每个子序列的长度为10，每个时间步的特征数为1
x = keras.Input(shape=input_shape)
y = keras.layers.Lambda(lambda x: x[:, -1, :])(x) # 将每个子序列的最后一个时间步的特征作为该子序列的输出
y = keras.layers.Reshape((1, 1))(y)

# 定义模型
model = transformer_model(
    input_shape=input_shape,
    num_layers=2,
    d_model=32,
    num_heads=4,
    dff=64,
    dropout_rate=0.2,
)

# 编译模型
model.compile(loss="mse", optimizer="adam")

# 训练模型
history = model.fit(x, y, epochs=100, batch_size=32)

我们首先定义了一个padding_mask张量，它是一个与输入张量形状相同的张量，其中每个元素的值为0或1，表示该位置是否是填充位置（如果是填充位置，则对应的值为1）。然后，我们使用Lambda层将输入张量转换为一个掩码张量，其中填充位置的值为0，非填充位置的值为1。接下来，我们使用Masking层将掩码张量应用于输入张量，从而对填充位置进行掩码。最后，我们在第一个Transformer层中传递掩码张量，以便模型能够在训练和推理过程中正确地使用掩码。

在Transformer模型中，Multi-Head Attention是一个重要的组件，它允许模型同时关注输入序列的不同位置，并且可以学习输入序列中不同位置之间的关系。在Multi-Head Attention中，我们需要传递两个参数：一个查询序列和一个键值对序列。在实际实现中，我们通常使用同一个输入序列来构建这两个序列，因此在Keras中实现时，传递的参数为[x,x]，其中x是输入序列。

具体来说，Multi-Head Attention包括三个线性变换，分别是查询、键和值的线性变换。在Keras中，我们可以使用一个全连接层（Dense层）来实现这个线性变换：

query = layers.Dense(d_model)(x)
key = layers.Dense(d_model)(x)
value = layers.Dense(d_model)(x)

在实现这个线性变换后，我们将query、key和value分别传递给Multi-Head Attention层。在Keras中，我们使用MultiHeadAttention层来实现Multi-Head Attention。这个层接受两个输入：一个查询序列和一个键值对序列。在实际应用中，我们通常使用同一个输入序列来构建这两个序列，因此在Keras中，我们需要将x复制一份，作为查询序列和键值对序列的输入：

attention_output = layers.MultiHeadAttention(num_heads=num_heads, key_dim=d_model)([query, key, value])

其中，num_heads表示头的数量，key_dim表示键和值的维度。在这个层中，我们将查询、键和值分别按头的数量进行划分，并对每个头进行独立的注意力计算。最终，我们将每个头的输出连接起来，形成最终的输出。注意，这个输出的形状与输入序列的形状相同。

因此，在Keras中实现Multi-Head Attention时，我们需要将同一个输入序列复制一份，作为查询序列和键值对序列的输入，然后将它们传递给MultiHeadAttention层。因此，输入的参数是[x,x]，其中x是输入序列。