提高 I/O bound tensorflow 训练速度答案

【问题标题】：Increase I/O bound tensorflow training speed提高 I/O bound tensorflow 训练速度
【发布时间】：2020-06-03 16:50:43
【问题描述】：

我面临着提高点云对象检测算法的 Tensorflow 实现的训练速度/效率的问题。

输入数据是 [8000, 100, 9] float32 张量，每个样本的大小约为 27MB。在批量大小为 5 时，数据加载成为训练中的瓶颈，因为在数据到达之前，GPU 使用率大部分时间都是 0%。

我尝试了以下方法来提高数据加载速度。

在 tf.Dataset .map API 中使用 num_parallel_calls，并使用多个线程来读取这个大张量。问题是.map 包装了一个py_fun，它受全局解释器锁的约束，因此多线程不会提高 I/O 效率。
使用 tf.Dataset .interleave API。由于它也是基于多线程的，所以和 2 有同样的问题。
使用 TFRecord 格式。这甚至比方法 1 和 2 还要慢。可能是 TFRecord 会将张量转换为 numpy，然后将 numpy 序列化为字节，然后将这些字节包装到 tensorflow 结构并写入磁盘。根据tf.convert_to_tensor() 测量的数据，Numpy to Tensor 需要很长时间。

任何关于如何前进的建议都会有所帮助。谢谢！

跟进cmets

我使用的是慢速磁盘吗？数据存储在已安装的磁盘上。可能是一个原因。
数据可以放入 GPU 内存吗？很不幸的是，不行。大约有 70,000 个样本。我尝试将一个小数据集缓存到 RAM 中，GPU 使用率为 30%~40%，这可能是这个特定网络的最高期望。

【问题讨论】：

我假设您正在从磁盘读取这些数据。你在使用慢速磁盘吗？后续问题，实际数据集有多大？因为如果它适合您的 GPU，您可以使用 Rapids 库，将其读入 GPU 并完全在 GPU 上进行处理。我不知道您可以使用哪种硬件，但该库是为此类吞吐量问题而设计的。特别是如果您可以访问 GPU 网格。
Apache spark 可用于此目的，towardsdatascience.com/…

【解决方案1】：

一些想法：

您应该使用 1,2 和 3 的组合。如果您将文件保存为 TFRecords，则可以并行读取它们，这就是它们的设计目的。然后，您将能够使用num_parallel_calls 和interleave，因为这样您就不必包装py_func。
.map 不必包装.py_func，例如，您可以使用tf.keras.utils.get_file。这样您也可以避免使用py_func 并有效地使用num_parallel_calls。我仍然推荐使用TFRecords，它们是为这个用例设计的。
另一种选择是使用 SSD 而不是硬盘来存储数据。
您还可以查看tf.Dataset API 的.cache 函数。也许您可以尝试加载数据的一个随机子集，在其上训练多个 epoch，然后同时获取另一个数据子集（使用tf.prefetch），然后在其上训练多个 epoch，依此类推。这个想法更像是一个长镜头，因为它可能会影响性能，但它可能适用于您的情况。

【讨论】：