扩展 tf.io.gfile.GFile 超过 100MB/s 的吞吐量答案

【问题标题】：Scaling tf.io.gfile.GFile over 100MB/s throughput扩展 tf.io.gfile.GFile 超过 100MB/s 的吞吐量
【发布时间】：2020-01-10 18:33:44
【问题描述】：

为了在训练期间充分利用 GPU，我需要能够向 GPU 提供大约 250 MB/s 的原始数据（数据不可压缩）。我正在通过快速网络访问数据，该网络可以毫无问题地提供超过 2GB/秒的数据。 Python 的 GIL 使得在不对训练循环产生负面影响的情况下，将这些速度融入到运行 TensorFlow 的同一进程中变得相当困难。 Python 3.8 的共享内存可能会缓解这种情况，但 Tensorflow 尚不支持。

所以我使用tf.io.gfile.GFile 通过网络读取数据（数据存储在高带宽 S3 兼容接口上）。 GFile 的价值在于它不参与 GIL，因此可以很好地与训练循环配合使用。为了实现高吞吐量，需要对网络 IO 进行显着的并行化。

不过，我似乎只能通过这种方法获得大约 75-100 MB/秒的速度。

我已经为两种方法计时：

创建一个tf.data.Dataset 并使用tf.data.Dataset.map(mymapfunc, num_parallel_calls=50)（我尝试了许多num_parallel_calls 的值，包括AUTOTUNE）。
创建一个使用tf.io.gfile.GFile 读取数据的函数，并在concurrent.futures.ThreadPoolExecutor 中使用多个线程简单地运行它，尝试线程数达到大约100 个（大约20 个以上没有任何改进，最终更多的线程会减慢它的速度）。

在这两种情况下，我的最高速度都是 75-100 MB/秒。

问题：

我想知道GFile 是否有理由达到可能更高的上限对其他人来说是显而易见的。

我也在做一个我应该验证的假设：tf.io.gfile.GFile 在 numpy 土地上运行，在上述两种情况下，我都在运行 GFile 操作来自 python 土地（在 tf.data.Dataset 的情况下，我正在使用 tf.py_function)。如果 GFile 旨在作为图形的一部分运行更有效地运营我不知道这一点，需要已更正。

【问题讨论】：

是的，在 S3 方面，管理员监控显示当我运行多个作业时，在我这边，我使用多个运行 asyncio 的 Python 进程将 1+GB/sec 基准测试到单个 pod以及针对 S3 接口的 500-1000 次并行操作。这是一个运行 Ceph 分布式文件系统的大型 kubernetes 集群。我非常确信瓶颈不在 S3 上，但如果其他人从 gfile 中看到高于 100MB/秒的速度，我很想听听。
但是，如果不将 GIL 完全锁定在主要 python 进程中，您就无法将 1GB/秒的数据从一个 python 进程移动到另一个进程，这会阻止 tensorflow 训练循环快速迭代，从而阻碍GPU ，所以在我可以尝试 python 3.8 和 SystemV 共享内存之前，多处理和多线程都不能很好地与 GPU 上的训练循环配合使用。这就是 gfile 的吸引力所在，它不会锁定 GIL。
当您说您尝试过Dataset.map 时，是不是就像DatasetReader(filename).map(...)？这试图通过一个来压缩所有数据。你看过Dataset.interleave吗，它更像FileNamesDataset().interleave(lambda f:DatasetReader(f))通过多个连接流式传输......曾经有一个“sloppy_interleave”选项可以提高性能，因为它可以跳过停滞的连接......但这似乎已经消失了tf2 可能隐藏在tf.data.Options.experimental_deterministic 中。
map 接收到一组要下载的文件段（文件名、偏移量、长度），大小为 5-30 MB，map 函数为每个段打开一个 gfile 并执行 seek 和read。针对 S3 接口打开 gfile 似乎没有任何明显的开销，因此我没有缓存该操作。我期待num_parallel_calls 是并行化的。由于我在线程池中看到了相同的行为，因此我并没有过多地质疑 map 操作。
“地图接收到一组要下载的文件段”听起来不错。 “大小为 5-30 MB”我听说每个文件 100 MB 的建议以获得最佳性能，但谁知道这是否会有所帮助。

标签： python-3.x tensorflow tensorflow-datasets

【解决方案1】：

对我来说这个问题的最终解决方案如下：

我无法让 GFile 以特别高的速度下载，boto3.Client.download_file 使用多个线程要快得多。
我已经启动了一个单独的进程 (multiprocessing) 来通过 boto3 处理 IO，它会在本地下载大型 (TB+) 数据集的连续滚动窗口。
然后另一个进程对本地文件进行采样并将 TF Records 文件写入 RAM 磁盘，然后在写入文件名时将文件名传递给 tf.data.TFRecordDataset。
有必要通过 TFRecords 文件将大规模数据流移交给 tensorflow，因为 Python 中的反序列化会锁定 GIL 并延迟训练循环（特别是 Python 3.8 的 System V 共享内存可能会解决这个问题，但 TF 2.2 is the first version expected to support Python 3.8 )。
RAM 磁盘避免了 IO 过载，但我仍然需要 NVMe 级别的 IO 性能来同时下载（100 MB/秒 IO）和读取原始数据文件并生成 tf 记录（200 MB/秒 IO）。李>
我处理不完全适合本地的 TB+ 数据集的最终解决方案是通过 boto3 在本地下载连续滚动的原始数据文件窗口，并从可用的本地文件中运行 2 个进程采样并将 TF 记录文件写入RAM 磁盘，然后在生成这些文件名时将它们交给tf.data.TFRecordDataset。组件可以以几乎完全解耦的方式编写，因此调试起来并不可怕。
使用此解决方案，GPU 保持 >90%。 IO 需要大约 300MB/秒的聚合读/写。内存使用量：RAM 磁盘中约 10GB，进程中约 5GB RAM。该模型以大约 200MB/秒的原始（不可压缩）数据为基础。

【讨论】：