TensorFlow for循环内Lambda层张量切片仅返回最后索引列问题求解

这是Python闭包的延迟绑定特性导致的：你在循环中定义的Lambda匿名函数没有立即捕获当前迭代的i值，等到模型实际推理执行函数时，循环已经跑完，i的值固定为最后一次迭代的结果R-1，所以所有切片都只会取张量的最后一列。

只需要给匿名函数添加默认参数，就能在每次迭代时把当前i的值绑定到函数内部，不需要改动其他逻辑：

R=2
inp = tf.keras.Input(shape = (R,), dtype=tf.int32)

SLICES = []
for i in range(R):
    # 增加i=i的默认参数，绑定当前迭代的i值
    slice_ = tf.keras.layers.Lambda(lambda a, i=i: a[:,i], name=f"slice_{i}", dtype=tf.int32)(inp)
    SLICES.append(slice_)

model = tf.keras.Model(inputs= inp, outputs = SLICES)

修改后输入测试数据就能得到预期的拆分结果：

[array([1, 3, 5], dtype=int32), array([2, 4, 6], dtype=int32)]

不需要手动循环写Lambda层，直接用tf.unstackAPI即可按列拆分张量，代码更简洁、执行效率更高，也不会遇到闭包绑定的坑：

R=2
inp = tf.keras.Input(shape = (R,), dtype=tf.int32)
# 直接按axis=1（列维度）拆分为R个张量
SLICES = tf.unstack(inp, num=R, axis=1)
model = tf.keras.Model(inputs= inp, outputs = SLICES)

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