Keras 3 中使用 tf.data.Dataset 和 steps_per_execution 的注意事项

在 Keras 3 框架中，当开发者尝试结合使用 tf.data.Dataset 数据管道和 steps_per_execution 参数（值大于32）时，会遇到一个常见的运行时错误。这个错误提示表明检测到了异常高数量的 tf.data.Iterator.get_next() 调用，暗示在禁用 AutoGraph 的情况下使用了不支持的迭代模式。

问题背景

Keras 3 默认禁用了 TensorFlow 的 AutoGraph 功能，这是为了优化性能而做出的设计选择。当开发者设置 steps_per_execution 参数大于32时，框架内部会尝试在一个训练步骤中执行多次迭代操作。由于 AutoGraph 被禁用，这种批量迭代操作会触发 TensorFlow 数据管道的保护机制，导致运行时错误。

技术原理

TensorFlow 数据管道（tf.data）实现了一个保护性启发式算法，当检测到单个函数中执行过多 get_next() 操作时会抛出警告。这个机制原本是为了防止开发者意外地在禁用 AutoGraph 的情况下使用 Python 原生的 for 循环迭代数据集。

在 Keras 的训练循环中，当 steps_per_execution 设置较大时，框架会生成一个包含多个步骤的函数。如果这个函数使用 Python 原生的 range 进行循环，且 AutoGraph 被禁用，就会触发上述保护机制。

解决方案

目前有几种可行的解决方案：

1. 启用 Eager 执行模式：通过在 model.compile() 中设置 run_eagerly=True 可以绕过这个问题，因为 Eager 模式会自动启用 AutoGraph。但这种方法会牺牲部分性能优势。

2. 修改训练函数实现：更底层的解决方案是修改 Keras 的 make_train_function 实现，将 Python 原生的 range 替换为 tf.range。这样即使 AutoGraph 被禁用，TensorFlow 也能正确地将循环转换为图操作。

3. 限制 steps_per_execution：作为临时解决方案，可以将 steps_per_execution 限制在32以内，避免触发保护机制。

最佳实践

对于追求性能的开发者，推荐采用第二种方案，即修改训练循环的实现方式。这种方案既保持了性能优势，又解决了兼容性问题。开发者可以继承 Keras 的 Trainer 类，重写 make_train_function 方法，确保使用 tf.range 而不是 Python 原生的 range。

这个问题的出现也提醒我们，在深度学习框架升级时，需要注意底层机制的变化可能带来的兼容性问题。Keras 3 为了提高性能所做的默认配置调整，可能会影响到一些特定的使用场景。