在 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 之间启用 p2p 通信时，单塔模型在 GPU 上保持变量更快，多塔模型也更快。要将变量固定到 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 队列运行器方法快得多。他们正在记录中。