如何将接收一维张量的tf.keras模型应用于任意形状张量？

针对你的场景（模型输出仅依赖对应位置的输入元素，无跨元素依赖），最简洁高效的方案是先将任意形状的输入张量展平为一维，通过模型处理后再恢复原始形状，完全不需要递归使用tf.map_fn。

具体实现步骤 #

1. 记录输入张量的原始形状，方便后续恢复；
2. 将输入张量展平为一维形状(-1,)，适配模型的输入要求；
3. 用模型处理展平后的张量；
4. 将处理后的结果重塑回原始形状。

代码示例 #

假设你的预训练模型名为elementwise_model，输入张量为任意形状的input_tensor：

import tensorflow as tf

# 记录输入的原始动态形状
original_shape = tf.shape(input_tensor)
# 展平为一维张量
flattened_input = tf.reshape(input_tensor, (-1,))
# 用模型处理展平后的输入
flattened_output = elementwise_model(flattened_input)
# 恢复为原始形状
output_tensor = tf.reshape(flattened_output, original_shape)

为什么这个方法可行？ #

因为你的模型满足元素独立性（每个位置的输出仅依赖对应位置的输入），展平操作只是改变了张量的维度组织方式，并没有改变元素本身的顺序和值，处理后恢复形状完全不会影响最终结果。

相比tf.map_fn，这种方法的优势在于：

• 代码更简洁，无需递归处理多维度；
• 执行效率更高，TensorFlow会对展平后的批量元素做整体优化，避免循环带来的开销；
• 兼容性强，无需修改原有模型结构，适用于任何满足元素独立性的模型。

额外优化：修改模型支持任意维度输入 #

如果你希望模型直接接受任意形状的输入，也可以调整模型的输入层为接受任意形状的张量。例如，将原模型的输入层从适配(batch_size,)形状的Input(shape=())改为Input(shape=None)，这样模型就能直接处理(dim1, dim2, ..., dimn)形状的输入：

# 重新构建支持任意维度的模型
def build_elementwise_model():
    inputs = tf.keras.layers.Input(shape=None, dtype=tf.float32)
    # 替换为你原模型的层结构，确保所有操作都是元素级的
    x = tf.keras.layers.Dense(16, activation='relu')(inputs)
    outputs = tf.keras.layers.Dense(1)(x)
    # 挤压最后一维（匹配原模型输出的(batch_size,)形状）
    outputs = tf.keras.layers.Lambda(lambda x: tf.squeeze(x, axis=-1))(outputs)
    return tf.keras.Model(inputs=inputs, outputs=outputs)

# 使用新模型直接处理任意形状输入
model = build_elementwise_model()
output_tensor = model(input_tensor)  # input_tensor可以是(3,4)、(5,6,7)等任意形状

注意：这种方法要求模型内的所有层都支持任意维度输入（如Dense、Activation等元素级操作均支持，而LSTM等序列相关层则不适用），这刚好符合你的场景要求。

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