在 TensorFlow 中使用 GPU 时的内存管理

Question

我对在 Tensorflow 中使用 GPU 有一些疑问。我正在学习卷积神经网络教程 here (tensorflow/models/image/cifar10/cifar10_train.py)。正如教程中一样，所有参数(例如权重)都在 CPU 内存中存储和更新，GPU 仅用于计算梯度或推理。

由于权重存储在 CPU 中，因此它们应该在每次迭代时同步，并且 GPU 似乎未得到充分利用(根据 nvidia-smi 大约为 60%)。在使用多个 GPU 的情况下，我知道权重应该存储在 CPU 内存中，以便在 GPU 之间同步。然而，为什么本教程即使在单个 GPU 中也将所有权重存储在 CPU 中？有没有办法在 GPU 内存中存储和更新它们？

在推理时，权重是否复制到 GPU 一次并重复使用？还是每次使用时都应该复制？

图像数据怎么样？这些数据似乎驻留在 GPU 中(不确定)。这些数据什么时候传输到GPU？当它们从磁盘加载时？或者什么时候 GPU 需要它们？

• 如果从磁盘加载后立即复制到 GPU，如果图像数据的大小太大而无法放入 GPU 内存，会​​发生什么情况？在这种情况下，有什么办法可以单独复制数据(比如预取)？
• 如果按需将它们复制到 GPU，是否有任何方法可以在 GPU 实际使用它们之前预取它们以避免空闲时间？

编辑:如果有任何方法可以检查发送/接收节点在 CPU 和 GPU 之间插入的位置(如 white paper 中)，将会很有帮助。

Answer 1

这些教程旨在展示 API，因此它们不会针对性能进行优化。对于单塔模型，在 GPU 上保持变量的速度更快；当您在 GPU 之间启用 p2p 通信时，对于多塔模型，保持变量的速度也更快。要将变量固定到 GPU，请使用与任何其他操作相同的 tf.device('/gpu:0') 方法。

如果启用分区图，您可以看到 GPU 之间的所有内存副本，即执行如下操作:

metadata = tf.RunMetadata()
sess.run(x, options=tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE,
                                  output_partition_graphs=True),
         run_metadata=metadata)

timeline = Timeline(metadata.step_stats)
with open("dynamic_stitch_gpu_profile.json", "w") as f:
    f.write(timeline.generate_chrome_trace_format())
with open("dynamic_stitch_gpu_profile.pbtxt", "w") as f:
    f.write(str(metadata))

请参阅此问题，了解使用此技术追踪副本的示例: https://github.com/tensorflow/tensorflow/issues/7251#issuecomment-277385212

要预取到 GPU，请参阅此 issue

添加了新的 stage_op 操作，允许预取到 GPU，并且比使用 Python 队列运行器方法要快得多。它们正在被记录下来。