TensorFlow tf.data.Dataset多输入尺寸不同报形状不匹配错误

报错现象 #

构建包含两类不同输入的TensorFlow模型时，两类输入形状分别为[9,10]和[8]，使用tf.data.Dataset迭代数据抛出如下错误：

tensorflow.python.framework.errors_impl.InvalidArgumentError: Shapes of all inputs must match: values[0].shape = [9,10] != values[1].shape = [8]
     [[{{node packed}}]] [Op:IteratorGetNext]

已知TensorFlow函数式API支持将不同尺寸的输入传入不同模型分支，官方函数式API指南中也有对应示例，但示例未使用tf.data.Dataset，无法直接参考实现。

问题背景 #

当前处理多站点时间序列预测任务，输入分为两类：

• 随时间变化的动态协变量
• 不随时间变化、但存在站点间差异的静态协变量
  任务目标是预测下一个时间步的数值。

首先使用timeseries_dataset_from_array工具，通过滑动窗口生成动态协变量与对应目标值的数据集，代码如下，该部分运行正常，可直接用于模型训练：

train_ds = tf.keras.preprocessing.timeseries_dataset_from_array(
    input_data, targets, sequence_length=window_size, batch_size=256)

后续需要额外引入时间序列对应站点的静态协变量：窗口输入数据中有一列单独存储站点ID，该列会在训练前移除。尝试提取对应站点的静态协变量，作为独立输入加入数据集，代码如下：

train_ds = train_ds.map(lambda x, y: (tf.py_function(attach_static_covariates, [x, idindex, colnames], [tf.float64, tf.float64]), y))
train_ds = train_ds.map(lambda x, y: ({'dynamic': x[0], 'static': x[1]}, y))

其中attach_static_covariates方法实现如下：

def attach_static_covariates(x, idindex, colnames):
    id = x[0, idindex].numpy()
    static_cov = static_df.iloc[int(id)]
    # 过滤掉已经完成用途的ID列
    x = tf.gather(x, [i for i in range(len(colnames)) if i != idindex])
    return (x, static_cov)

已完成验证：如果两个输入尺寸相同（比如返回(x, x)，将两份动态协变量分别传入模型两个分支），上述代码可以正常运行并完成训练。该问题并非模型定义不匹配导致，直接执行如下迭代代码时也会抛出完全相同的错误：

for feature_batch, label_batch in train_ds.take(1):
    print(feature_batch)

检索过官方文档、代码仓库等渠道，未找到同类问题的解决方案。不同尺寸多输入是非常普遍的使用场景，按理tf.data.Dataset应该支持该用法，但目前没有找到用tf.data.Dataset加载不同形状多输入的公开示例，需要对应的解决方案或参考示例。

解决方案 #

报错核心原因是tf.py_function默认不会自动推断输出张量的静态形状，后续map操作组装输入字典时，TensorFlow会因两个输出张量形状未知，尝试做对齐打包触发形状校验。
修复时只需要在py_function逻辑执行完成后，显式给两个输出张量设置确定的静态形状即可，不需要修改模型结构。
修改后的代码示例：

def attach_static_covariates(x, idindex, colnames):
    id = x[0, idindex].numpy()
    static_cov = static_df.iloc[int(id)].values.astype('float64')
    # 过滤ID列得到动态协变量
    dynamic_x = tf.gather(x, [i for i in range(len(colnames)) if i != idindex])
    return (dynamic_x, static_cov)

def parse_dataset(x, y):
    # 调用py_function获取两类输入
    dyn_input, stat_input = tf.py_function(
        attach_static_covariates, 
        [x, idindex, colnames], 
        [tf.float64, tf.float64]
    )
    # 显式设置两个张量的静态形状，batch维度设为None即可
    dyn_input.set_shape([None, 9, 10])
    stat_input.set_shape([None, 8])
    return ({'dynamic': dyn_input, 'static': stat_input}, y)

train_ds = train_ds.map(parse_dataset)

修改完成后可直接通过迭代验证：执行for feature_batch, label_batch in train_ds.take(1): print(feature_batch['dynamic'].shape, feature_batch['static'].shape)，会正常输出两个不同的张量形状，不再触发形状不匹配错误，生成的数据集可直接传入多输入函数式模型完成训练。

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