使用tf.data管道时Keras模型无法训练的原因排查

针对你遇到的「Numpy数组训练正常，tf.data管道训练精度卡基线」的问题，结合你提供的代码，以下是最可能的原因及对应的排查/修复方案：

1. to_categorical 混用Numpy操作导致标签异常 #

你在y_ds的map中使用了Keras的to_categorical，这个函数返回的是Numpy数组而非TensorFlow张量。在tf.data管道中混用Numpy操作会破坏计算图连贯性，可能导致标签类型、形状或数据传递出现隐性问题（比如张量被强制转换为非预期类型，或计算图外操作导致梯度无法正常传播）。

修复方案：替换为TensorFlow原生的tf.one_hot操作：

y_ds = (
    ds
   .skip(T - 1)
   .map(lambda s: tf.cast(s[-1] - 1, tf.int32))  # 确保标签为整数类型
   .map(lambda y: tf.one_hot(y, depth=3))
)

2. X与y的样本对齐验证 #

虽然理论上X_ds和y_ds的样本数量应该一致，但实际生成过程中可能因window/flat_map逻辑出现偏差。建议手动验证两者的样本数和对应关系：

# 打印样本数量
print(f"X_ds样本数: {len(list(X_ds.as_numpy_iterator()))}")
print(f"y_ds样本数: {len(list(y_ds.as_numpy_iterator()))}")

# 对比前3组样本是否与Numpy版本一致
for x_np, y_np in zip(X_ds.take(3).as_numpy_iterator(), y_ds.take(3).as_numpy_iterator()):
    print("X样本形状:", x_np.shape)
    print("y样本:", y_np)

若样本数不一致，检查window和skip参数是否匹配：window(T, shift=1, drop_remainder=True)生成的样本数应为总时间步 - T + 1，ds.skip(T-1)后的样本数也需等于该值，否则说明原始数据维度逻辑有误。

3. repeat 操作的冗余与错误 #

你的代码中使用了.repeat(n_epochs * size_batch)，这会导致数据集被重复远超训练所需次数，结合model.fit中的epochs=n_epochs，可能让模型在重复数据循环中无法有效学习。正确做法是让tf.data自动处理epoch循环，无需手动指定repeat次数：

Xy_ds = (
    tf.data.Dataset.zip(X_ds, y_ds)
    .shuffle(buffer_size=1000)  # 若Numpy版本有shuffle，此处必须添加，否则模型易因顺序数据难以收敛
    .batch(size_batch)
    .prefetch(tf.data.AUTOTUNE)
)

注意：若Numpy版本训练时对数据做了shuffle，tf.data版本必须添加.shuffle()步骤——顺序数据会让模型极易陷入基线精度。

4. 数据形状与类型的隐性差异 #

即使磁盘数据一致，tf.data管道中的张量形状/类型可能与Numpy版本存在细微差异：

• 检查X样本形状：确保tf.expand_dims(x, -1)后的形状与Numpy版本完全一致（比如均为(T, n_features, 1)）。
• 检查数据类型：打印X_ds.take(1).as_numpy_iterator()的dtype，确认与Numpy数组的dtype完全匹配（比如均为float32）。

5. 线程/并行优化的影响 #

虽然你尝试了禁用线程，但tf.data的prefetch和map并行设置仍可能影响数据传递。可以先简化管道，关闭所有并行优化验证：

# 简化管道，关闭并行
X_ds = (
    ds
    .map(lambda s: s[:n_features], num_parallel_calls=1)
    .window(T, shift=1, drop_remainder=True)
    .flat_map(lambda x: x.batch(T, drop_remainder=True), num_parallel_calls=1)
    .map(lambda x: tf.expand_dims(x, -1), num_parallel_calls=1)
)

Xy_ds = (
    tf.data.Dataset.zip(X_ds, y_ds)
    .batch(size_batch)
    .prefetch(1)  # 禁用AUTOTUNE，改为固定值1
)

若简化后训练恢复正常，再逐步开启并行优化，排查具体是哪个步骤导致的问题。

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