python-3.x - 扩展 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 兼容接口(interface)上)。 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/秒。

问题:

I'm wondering if there's a reason for GFile to hit an upper limit that is perhaps more obvious to someone else.

I'm also making an assumption I should validate: tf.io.gfile.GFile runs in numpy land, in both cases above I'm running GFile operations from python land (in the case of tf.data.Dataset I'm using tf.py_function). If GFile is meant to run as part of the graph operations more efficiently I'm unaware of this and need to be corrected.

Answer 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 Records 文件写入 RAM 磁盘，然后将这些文件名交给 tf.data.TFRecordDataset因为它们是生成的。组件可以以几乎完全解耦的方式编写，因此调试起来并不可怕。

使用此解决方案，GPU 保持 >90%。 IO 需要大约 300MB/秒的聚合读/写。内存使用量:RAM 磁盘中约 10GB，进程中约 5GB RAM。该模型以大约 200MB/秒的原始(不可压缩)数据为基础。