tf.data.Dataset.interleave() 与 map() 和 flat_map() 究竟有何不同？

Posted 2023-02-16

【中文标题】tf.data.Dataset.interleave() 与 map() 和 flat_map() 究竟有何不同？

【英文标题】：How exactly does tf.data.Dataset.interleave() differ from map() and flat_map()?

【发布时间】：2021-06-20 23:49:48

【问题描述】：

我目前的理解是：

不同的 map_func：interleave 和 flat_map 都期望“将数据集元素映射到 数据集 的函数”。相比之下，map 需要“将数据集元素映射到另一个 数据集元素 的函数”。

参数：interleave 和 map 都提供参数 num_parallel_calls，而 flat_map 不提供。此外，interleave 提供了这些神奇的参数 block_length 和 cycle_length。对于 cycle_length=1，文档说明 interleave 和 flat_map 的输出是相等的。

最后，我见过data loading pipelines without interleave 和ones with interleave。任何建议何时使用 interleave 与 map 或 flat_map 将不胜感激

---

//编辑：我确实看到交错的价值，如果我们从不同的数据集开始，例如在the code below

  files = tf.data.Dataset.list_files("/path/to/dataset/train-*.tfrecord")
  dataset = files.interleave(tf.data.TFRecordDataset)

但是，在如下场景中使用interleave 比使用map 有什么好处吗？

files = tf.data.Dataset.list_files("/path/to/dataset/train-*.png")
dataset = files.map(load_img, num_parallel_calls=tf.data.AUTOTUNE)

【参考方案1】：

编辑：

map 不能也用于并行化 I/O 吗？

确实，您可以使用map 函数从目录中读取图像和标签。假设这种情况：

list_ds = tf.data.Dataset.list_files(my_path)

def process_path(path):
 ### get label here etc. Images need to be decoded
 return tf.io.read_file(path), label

new_ds = list_ds.map(process_path,num_parallel_calls=tf.data.experimental.AUTOTUNE)

注意，现在它是多线程的，因为已经设置了num_parallel_calls。

interlave()功能的优势：

假设您有一个数据集 使用cycle_length，您可以从数据集中取出许多元素，即5，然后从数据集中取出5 个元素，然后可以应用map_func。 之后，从新生成的对象中获取数据集对象，每次block_length条数据。

换句话说，interleave() 函数 c在应用 map_func() 的同时遍历您的数据集。此外，它可以同时处理许多数据集或数据文件。例如from the docs:

  dataset = dataset.interleave(lambda x:
    tf.data.TextLineDataset(x).map(parse_fn, num_parallel_calls=1),
    cycle_length=4, block_length=16)

但是，使用 interleave over map 有什么好处吗？ 像下面这样的场景？

interleave() 和 map() 看起来有点相似，但它们的用例并不相同。如果您想在应用一些映射的同时读取数据集，interleave() 是您的超级英雄。您的图像可能需要在读取时进行解码。首先读取所有内容，在处理大型数据集时解码可能效率低下。在您提供的代码 sn-p 中，AFAIK，带有tf.data.TFRecordDataset 的代码应该更快。

TL;DR interleave() 通过交错 I/O 操作来并行化数据加载步骤以读取文件。

map() 将对数据集的内容进行数据预处理。

所以你可以这样做：

ds = train_file.interleave(lambda x: tf.data.Dataset.list_files(directory_here).map(func,
                            num_parallel_calls=tf.data.experimental.AUTOTUNE)

tf.data.experimental.AUTOTUNE 将决定缓冲区大小、CPU 能力以及 I/O 操作的并行级别。换句话说，AUTOTUNE 将在运行时动态处理关卡。

num_parallel_calls 参数产生多个线程以利用多个内核来并行化任务。有了这个，您可以并行加载多个数据集，减少等待文件打开的时间；因为interleave 也可以接受参数num_parallel_calls。图片拍摄于from docs。

在图像中，有 4 个重叠的数据集，由参数 cycle_length 确定，因此在本例中为 cycle_length = 4。

---

FLAT_MAP：跨数据集映射函数并将结果展平。如果您想确保订单保持不变，您可以使用它。而且它不以num_parallel_calls 作为参数。更多详情请参考docs。

地图： map 函数将分别对数据集的每个元素执行选定的函数。显然，随着您应用越来越多的操作，大型数据集上的数据转换可能会很昂贵。关键是，如果没有充分利用 CPU，它可能会更耗时。但是我们可以使用parallelism APIs:

num_of_cores = multiprocessing.cpu_count() # num of available cpu cores
mapped_data = data.map(function, num_parallel_calls = num_of_cores)

对于 cycle_length=1，文档指出 interleave 和 flat_map 相等

cycle_length --> 将同时处理的输入元素的数量。设置为1时，会一一处理。

INTERLEAVE：像map这样的变换操作可以并行化。

利用map的并行性，在顶部，CPU试图在转换中实现并行化，但是从磁盘中提取数据会导致开销。

此外，一旦将原始字节读入内存，可能还需要将函数映射到数据，这当然需要额外的计算。像解密数据等。需要并行化各种数据提取开销的影响，以便通过交错每个数据集的内容来减轻这种影响。

因此，在读取数据集时，您希望最大化：

图片来源：deeplearning.ai

【讨论】：

嗨@Frightera，我真的很感谢你花时间回答和投票。但是，不幸的是，我很难从中提取 map() 和 interleave() 之间的关键区别:(。一个是否只并行化 CPU 任务而另一个也并行化 I/O 任务？它与进程和线程有关吗？另外，你能引用一篇文章，你从中提取了图像吗？

我编辑了答案，看看现在是否清楚。我实际上是从视频中截取的图像，该视频可在 coursera 上找到。