TensorFlow中dataset.map在图模式下无限运行不返回问题咨询

问题重现 #

以下代码在即时执行模式下正常运行，但取消@tf.function注释后，图模式执行时会陷入无限运行：

import tensorflow as tf
dataset = tf.data.Dataset.from_tensor_slices(tf.constant(range(3)))
res_map = dataset.map(
    lambda x: x*2
)

# @tf.function # hangs in graph mode
def outer():
    return tf.convert_to_tensor(list(res_map))
res = outer()
print(res)

即时执行模式返回预期结果：

tf.Tensor([0 2 4], shape=(3,), dtype=int32)

问题原因 #

这是代码逻辑的根本性问题，而非TensorFlow的缺陷。在@tf.function修饰的图模式函数中，使用Python原生的list()去迭代tf.data.Dataset会触发无限循环：

• 图模式下，TensorFlow的图引擎负责处理Dataset的迭代逻辑，而Python的list()属于Python运行时操作，两者执行模型不兼容。
• list()会尝试不断从Dataset中获取元素，但图模式下无法正确识别Dataset的终止信号，导致持续循环。

修复方案 #

必须使用TensorFlow原生的Dataset操作完成转换，避免在图函数中混用Python迭代逻辑。以下是两种可行的修改方式：

方式1：使用batch合并所有元素 #

import tensorflow as tf
dataset = tf.data.Dataset.from_tensor_slices(tf.constant(range(3)))
res_map = dataset.map(lambda x: x*2)

@tf.function
def outer():
    # 用batch(-1)将所有元素合并为一个批次，再获取单个元素张量
    return res_map.batch(-1).get_single_element()

res = outer()
print(res)

方式2：使用reduce聚合元素 #

import tensorflow as tf
dataset = tf.data.Dataset.from_tensor_slices(tf.constant(range(3)))
res_map = dataset.map(lambda x: x*2)

@tf.function
def outer():
    # 通过reduce逐步拼接所有元素为一个张量
    return res_map.reduce(
        initial_state=tf.constant([], dtype=tf.int32),
        reduce_func=lambda acc, x: tf.concat([acc, [x]], axis=0)
    )

res = outer()
print(res)

两种方式在图模式下都能正常返回预期结果：tf.Tensor([0 2 4], shape=(3,), dtype=int32)

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