先导入Pandas后调用TensorFlow Dataset的next()陷入无限循环的原因

问题根源 #

这是Pandas与TensorFlow在Python运行时层面的迭代器协议冲突导致的：

• Pandas 2.2.x版本在导入时会修改底层迭代器相关逻辑，干扰了Python标准迭代器协议的正常执行流程。
• TensorFlow的tf.data.Dataset迭代器依赖标准迭代器协议实现，当Pandas先被导入后，take(1)生成的数据集迭代器无法正确识别终止条件，调用next()时会不断尝试从生成器获取元素，最终陷入无限循环。

复现代码 #

import pandas as pd  
import tensorflow as tf

# 模拟数据集生成器
def generator():
  for i in range(5):
    yield (tf.random.uniform(shape=(2, 7, 11)), tf.random.uniform(shape=(2, 1, 1)))

original_dataset = tf.data.Dataset.from_generator(
    generator,
    output_signature=(
        tf.TensorSpec(shape=(None, 7, 11), dtype=tf.float32),
        tf.TensorSpec(shape=(None, 1, 1), dtype=tf.float32)
    )
)

# 获取首个元素的数据集
first_element_ds = original_dataset.take(1)
iterator = iter(first_element_ds)

# 此处调用next()会陷入无限循环
first_element_tuple = next(iterator)

tensor1, tensor2 = first_element_tuple
print("Tensor 1:", tensor1.numpy().shape)
print("Tensor 2:", tensor2.numpy().shape)

环境信息 #

• Pandas版本：2.2.2
• TensorFlow版本：2.17.0
• Python版本：3.11.9
• NumPy版本：1.26.4
• Jupyter Core版本：5.7.2
• 操作系统：macOS 13.7.4
• 虚拟环境：conda 24.11.0，cPython 3.11.5

解决方案 #

方案1：调整导入顺序 #

先导入TensorFlow，再导入Pandas，避免Pandas干扰TensorFlow的迭代器逻辑：

import tensorflow as tf
import pandas as pd

# 后续代码保持不变

方案2：使用TensorFlow原生迭代器 #

改用as_numpy_iterator()生成不受干扰的迭代器：

first_element_ds = original_dataset.take(1)
iterator = first_element_ds.as_numpy_iterator()
first_element_tuple = next(iterator)

方案3：直接获取单个元素 #

如果确定数据集只有一个元素，使用get_single_element()方法直接获取，无需手动迭代：

first_element_tuple = original_dataset.take(1).get_single_element()

内容的提问来源于stack exchange，提问作者J.E.K