将 tf.io.gfile.GFile 扩展到 100MB/s 以上的吞吐量

Question

为了在训练期间充分利用 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，并使用 a 中的多个线程简单地运行它concurrent.futures.ThreadPoolExecutor，尝试将线程数增加到大约 100（超过大约 20 后没有任何改进，最终更多的线程会减慢速度）。

在这两种情况下，我的速度都达到了 75-100 MB/秒。

问题：

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

我还做出了一个应该验证的假设：tf.io.gfile.GFile 在 numpy land 中运行，在上面的两种情况下，我都从 python land 运行操作（在我使用的 GFile情况下）。如果 GFile 旨在作为图形操作的一部分更有效地运行，我不知道这一点并且需要更正。tf.data.Datasettf.py_function

Answer 1

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

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