使用TensorFlow的Dataset API和多GPU培训

Question

使用Tensorflow的新数据集API进行多GPU培训（来自TFRecords格式）似乎比在单个GPU（1对4特斯拉K80）​​上运行要慢得多（慢1/4）。

查看nvidia-smi的输出看来，使用4个GPU只会导致gpu利用率各占15％左右，而单个GPU则约为45％。

从磁盘加载数据（tfrecords-format）是否会导致训练速度瓶颈？使用常规feed-dicts，整个数据集加载到内存中的速度比使用数据集API要快得多。

Answer 1

您的网络似乎受到以下限制：

1. 来自光盘的IO，如上一段所述 如果您在读取TFRecords时启动数据集，那么它将从光盘读取;相反，您可以将它们读入列表/字典，并从范围序列开始。例如。

tf.data.Dataset()\ .range(your_data_size)\ .prefetch(20)\ .shuffle(buffer_size=20)\ .map(lambda i: your_loaded_list[i], num_parallel_calls=8)

1. 重型前/后处理，如第2段中所述，单GPU利用率为45％;如果那时你已经将数据预先加载到内存中，那么它会建议你的网络在“主要”计算机构之外做出努力。

首先，您可以尝试检查是否使用多线程与上面的地图调用有关;还会削减一些tf.summary操作，这些操作可能会反馈大量不必要的数据，这些数据会限制你的带宽/写入光盘。

希望这有帮助。