tf.data.Dataset.zip替代方案：如何生成两数据集全量元素组合

默认的tf.data.Dataset.zip逻辑和Python内置zip一致，只会对位匹配两个数据集相同索引的元素，无法生成所有元素的两两组合。要实现笛卡尔积形式的全量配对，直接用flat_map嵌套map即可，不需要额外引入其他依赖。

实现代码 #

import tensorflow as tf

a = tf.data.Dataset.range(1, 4)  # 数据集元素: [1, 2, 3]
b = tf.data.Dataset.range(4, 7)  # 数据集元素: [4, 5, 6]

# 生成全量配对数据集
full_pair_ds = a.flat_map(lambda a_elem: b.map(lambda b_elem: (a_elem, b_elem)))

# 打印结果验证
print(list(full_pair_ds.as_numpy_iterator()))
# 输出: [(1, 4), (1, 5), (1, 6), (2, 4), (2, 5), (2, 6), (3, 4), (3, 5), (3, 6)]

实现逻辑说明 #

• 外层用flat_map遍历第一个数据集a的每一个元素
• 对每一个拿到的a中元素，内层遍历第二个数据集b的所有元素，和当前a的元素组装成元组
• flat_map会自动把内层生成的所有子数据集展平，最终输出包含所有配对结果的单数据集

注意：该方案生成的数据集总长度是两个原始数据集长度的乘积，当两个数据集体量都很大时，总数据量会快速增长，要提前评估开销。

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