提高IO约束张量流训练速度

Question

我面临的是一个Tensorflow实现的点云对象检测算法的训练速度效率提升问题。

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

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

1. 使用 num_parallel_calls 在tf.Dataset中 .map API，并使用多个线程来读取这个大张量。问题是 .map 包裹 py_fun 的全局解释器锁，因此多线程并不能提高IO效率。
2. 使用tf.Dataset .interleave API。因为它也是基于多线程的，所以它的问题和2一样。
3. 使用TFRecord格式。这比方法1和2更慢。可能是TFRecord会把tensor转换成numpy，然后把numpy序列化成字节，再把这个字节包装成tensorflow结构，然后写入磁盘。Numpy转Tensor对我的数据来说需要很长的时间，以我的数据为例，我的数据是由 tf.convert_to_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. 你也可以研究一下 .cache 的作用 tf.Dataset API。也许你可以尝试加载一个随机的数据子集，在这个子集上训练多个eopchs，然后在这段时间内获取另一个数据子集（使用 tf.prefetch)，然后在此基础上训练多个纪元，以此类推。这个想法比较冒险，因为它可能会影响性能，但在你的情况下可能会有用。