python - 提高 I/O 绑定(bind) tensorflow 训练速度

Question

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

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

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

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

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

跟进评论

1. 我使用的是慢速磁盘吗？数据存储在已安装的磁盘上。可能是一个原因。
2. 数据是否适合 GPU 内存？抱歉不行。大约有 70,000 个样本。我尝试将一个小数据集缓存到 RAM 中，GPU 利用率为 30%~40%，这可能是对该特定网络的最高期望值。

Answer 1

一些想法:

1. 您应该使用 1,2 和 3 的组合。如果您将文件保存为 TFRecords，您可以并行读取它们，这就是它们的设计目的。然后，您将能够使用 num_parallel_calls 和 interleave，因为这样您就不必包装 py_func。

2. .map 不必包装 .py_func，例如您可以使用 tf.keras.utils.get_file。这样您也可以避免使用 py_func 并有效地使用 num_parallel_calls。我仍然推荐使用 TFRecords，它们是为这个用例设计的。

3. 另一种选择是使用 SSD 代替硬盘来存储数据。

4. 您还可以查看 tf.Dataset API 的 .cache 函数。也许您可以尝试加载数据的随机子集，在其上训练多个 eopch，然后同时获取数据的另一个子集(使用 tf.prefetch)，然后在上训练多个 epoch那个，等等。这个想法更像是一个远景，因为它可能会影响性能，但它可能适用于您的情况。