TensorFlow 2.16中Keras Lambda层触发ValueError异常求助

问题根源 #

TensorFlow 2.16对Lambda层的参数传递规则做了严格限制：只有输入张量能作为位置参数传入Lambda层的调用方法，非张量类型的参数（比如你代码里的整数256）必须以关键字参数形式传递。而2.15版本能正常运行，是因为旧版本对该规则的检查没有那么严格。

你的代码问题出在这里：

embeddings *= tf.keras.layers.Lambda(
    lambda model: tf.math.sqrt(tf.cast(model, tf.float32))
)(model)

这里的model是一个整数（即embedding的维度，建议改名为d_model避免和模型对象混淆），属于非张量参数，直接作为位置参数传给Lambda层的调用，触发了新版本的参数检查报错。

修复方案 #

方案一：直接计算（推荐） #

完全不需要用Lambda层，直接在外部完成平方根计算，代码更简洁高效：

inputs = tf.keras.Input(shape=(None,), name="inputs")
padding_mask = tf.keras.Input(shape=(1, 1, None), name="padding_mask")

# 把变量名model改成d_model，避免命名混淆
d_model = 256
embeddings = tf.keras.layers.Embedding(vocabulary_size, d_model)(inputs)
embeddings *= tf.math.sqrt(tf.cast(d_model, tf.float32))

方案二：保留Lambda层（按规则传参） #

如果一定要用Lambda层，需要通过arguments参数把非张量值传入，同时调整Lambda内部的函数逻辑：

inputs = tf.keras.Input(shape=(None,), name="inputs")
padding_mask = tf.keras.Input(shape=(1, 1, None), name="padding_mask")

d_model = 256
embeddings = tf.keras.layers.Embedding(vocabulary_size, d_model)(inputs)
embeddings *= tf.keras.layers.Lambda(
    lambda x, dim: tf.math.sqrt(tf.cast(dim, tf.float32)),
    arguments={"dim": d_model}
)(embeddings)

这里把整数参数d_model通过arguments传入Lambda层，Lambda内部函数的第一个参数是输入张量x，第二个是传入的关键字参数dim。

内容的提问来源于stack exchange，提问作者Martin