TensorFlow在.map中用py_function取张量值导致迭代器生成过慢如何解决

性能差异原因 #

• 上下文切换开销：tf.py_function 会在TensorFlow的C++运行时和Python解释器之间频繁切换执行上下文，每处理一个/一批样本都要触发一次切换，数据集规模大时累计开销极高。
• 失去图优化能力：自定义的编码函数用了Python层面的for循环、numpy运算，无法享受TensorFlow内置算子的图优化、GPU加速、并行计算等能力，运算效率远低于原生tf.one_hot。
• 数据拷贝开销：tf.keras.backend.get_value 会把张量从设备（GPU/TPU）拷贝到CPU内存转成numpy数组，高频调用时拷贝成本会被大幅放大。

map中处理张量值的方案 #

不需要特意将张量转为numpy值处理，几乎所有自定义逻辑都可以用TensorFlow原生算子在图内实现，完全可以避免使用tf.py_function。
以你当前的自定义独热编码需求为例，即使分类不是从0开始的连续值，也可以用原生算子实现：

# 自定义分类列表转为TF常量
CUSTOM_CATEGORIES = tf.constant([1,2,3,4,5,6,7,8,9,10], dtype=tf.int32)

def custom_one_hot(target):
    # 直接用张量运算实现匹配逻辑，无Python层面的循环和数据拷贝
    return tf.cast(target == CUSTOM_CATEGORIES, tf.float32)

# 映射时直接调用，性能和原生tf.one_hot一致
ds = ds.map(lambda input, target: (input, tf.reshape(custom_one_hot(target), (10,))),
            num_parallel_calls=tf.data.AUTOTUNE)

如果确实存在必须在Python层面处理张量的场景，可以通过以下方式优化性能：

1. 先对数据集做分批（batch）操作，再对整批数据应用tf.py_function，大幅减少上下文切换次数。
2. 调用map时设置num_parallel_calls=tf.data.AUTOTUNE开启并行处理，最后调用ds.prefetch(tf.data.AUTOTUNE)预加载数据，缓解性能瓶颈。

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