TensorFlow 2.15 不提供官方 Transformer 层,需手动组合 MultiHeadAttention、LayerNormalization 和前馈网络;key_dim 必须显式指定,mask 要通过 attention_mask 参数传入,位置编码需预计算并确保 shape 对齐,Dropout 必须传 training 参数。

TensorFlow 2.15 默认不提供现成的 Transformer 类,也没有 tf.keras.layers.Transformer —— 这个 API 从未在官方 stable 版本中存在过。想用原生 TF 搭建 Transformer,得自己组合 MultiHeadAttention、LayerNormalization 和前馈网络。
为什么不能直接 import Transformer?
很多人搜到旧教程里写的 from tensorflow.keras.layers import Transformer,实际在 TF 2.15 中会报 AttributeError: module 'tensorflow.keras.layers' has no attribute 'Transformer'。这是因为在 2021 年后,TF 官方移除了实验性 Transformer 层(曾短暂存在于 tf-nightly 的 tf.keras.layers.experimental),此后再未回归 stable 分支。
当前唯一可靠的路径是:手动拼装标准 Transformer 编码器块。
如何手搭一个可用的 EncoderLayer?
核心是复用 TF 内置的注意力与归一化组件,注意输入输出 shape 对齐和残差连接写法:
class EncoderLayer(tf.keras.layers.Layer):
def __init__(self, d_model, num_heads, dff, rate=0.1):
super().__init__()
self.mha = tf.keras.layers.MultiHeadAttention(
num_heads=num_heads, key_dim=d_model // num_heads)
self.ffn = tf.keras.Sequential([
tf.keras.layers.Dense(dff, activation='relu'),
tf.keras.layers.Dense(d_model)
])
self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
self.dropout1 = tf.keras.layers.Dropout(rate)
self.dropout2 = tf.keras.layers.Dropout(rate)
<pre class="brush:php;toolbar:false;">def call(self, x, training, mask=None):
# 注意:MultiHeadAttention 的 query/key/value 都来自 x
attn_output = self.mha(query=x, key=x, value=x, attention_mask=mask)
attn_output = self.dropout1(attn_output, training=training)
out1 = self.layernorm1(x + attn_output) # 残差 + norm
ffn_output = self.ffn(out1)
ffn_output = self.dropout2(ffn_output, training=training)
return self.layernorm2(out1 + ffn_output) # 残差 + norm-
key_dim必须显式指定(TF 2.15 不再自动推导),否则报ValueError: key_dim must be provided -
attention_mask要传给call()的mask参数,不能塞进**kwargs—— 否则掩码失效 - 所有
Dropout层必须显式传入training=training,否则训练/推理行为不一致
如何处理位置编码和输入嵌入?
TF 没有内置 PositionalEncoding 层,需手写并确保是不可训练的(或设为 trainable=False):
def positional_encoding(position, d_model):
angle_rads = np.arange(position)[:, np.newaxis] / np.power(
10000, (2 * (np.arange(d_model)[np.newaxis, :] // 2)) / np.float32(d_model))
angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])
angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])
pos_encoding = angle_rads[np.newaxis, ...]
return tf.cast(pos_encoding, dtype=tf.float32)
<h1>在模型 build 中调用:</h1><p>pos_encoding = positional_encoding(max_len, d_model)
x = x + pos_encoding[:, :tf.shape(x)[1], :]
</p>
- 别用
tf.range动态生成位置编码——会导致 SavedModel 导出失败;必须预计算或封装为tf.keras.layers.Lambda并设trainable=False - 嵌入层输出 shape 必须是
(batch, seq_len, d_model),否则和位置编码相加会广播错误 - 如果用
tf.keras.layers.Embedding做词嵌入,记得把input_length设对,否则后续 mask 推导出错
训练时 mask 怎么传才有效?
Transformer 的 padding mask 必须以 attention_mask 形式传入 MultiHeadAttention,且 shape 为 (batch, 1, 1, seq_len) 或 (batch, seq_len)(TF 会自动 broadcast)。常见错误:
- 用
tf.keras.layers.Masking层——它只影响 RNN/LSTM,对MultiHeadAttention无效 - 把 mask 当成普通张量加在输入上——这不会屏蔽 attention score,只会污染 embedding
- 在
EncoderLayer.call()里漏掉mask参数转发,导致下层无法接收
正确做法是在顶层模型中构造 mask:
def create_padding_mask(seq):
seq = tf.cast(tf.math.equal(seq, 0), tf.float32)
return seq[:, tf.newaxis, tf.newaxis, :] # (batch, 1, 1, seq_len)
<h1>然后在 call 中:</h1><p>mask = create_padding_mask(inp)
x = self.encoder_layer(x, training=training, mask=mask)
</p>
真正麻烦的不是写几层,而是 mask 传播链、shape 对齐、以及 SavedModel 兼容性——哪怕一个 tf.range 写错,导出时就静默失败。
Python免费学习笔记(深入):立即使用
在学习笔记中,你将探索 Python 的核心概念和高级技巧!











