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

我对在Tensorflow中使用GPU有些怀疑。我在学习卷积神经网络教程（tensorflow/models/image/cifar10/cifar10_train.py）。在本教程中，所有参数（例如权重）都存储并更新在CPU内存中，GPU仅用于计算梯度或推断

由于权重存储在CPU中，因此每次迭代都应同步权重，而且GPU似乎没有得到充分利用（根据

nvidia smi

，大约60%）。在使用多个GPU的情况下，我知道权重应该存储在CPU内存中，以便在GPU之间同步。然而，为什么本教程将所有权重存储在CPU中，甚至存储在单个GPU中？有没有办法在GPU内存中存储和更新它们

在推理的情况下，权重是否复制到GPU一次并重新使用？还是每次使用时都要复制

图像数据呢？这些数据似乎驻留在GPU中（不确定）。该数据何时传输到GPU？何时从磁盘加载它们？或者当GPU中需要它们时

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

---

编辑：如果有任何方法可以检查发送/接收节点插入CPU和GPU之间的位置（如中所示），这将非常有用。

这些教程旨在展示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))

有关使用此技术跟踪副本的示例，请参阅本期：

有关GPU的预取，请参见此

---

添加了新的

stage_op

op，允许对GPU进行预取，并且比使用Python队列运行器方法要快得多。它们正在被记录。

您可以使用队列将数据提取到GPU中。当计算进行时，它将被复制，并且有线程代码来保持它的完整性。p2p通信意味着类似GPUDirect的东西，对吗？那么如何在TensorFlow中使用它呢？如果我的GPU支持，TF是否自动使用GPUDirect？它通常意味着GPU位于同一PCI根复合体上。它会在可用时自动使用，当您第一次在时间线中启动GPU时，您将看到打印的p2p矩阵，每行的含义是什么？我可以看到带有/gpu:0/memcpy Compute]或/gpu:0/stream:XX Compute或/job:localhost/replica:0/task:0/cpu:0 Compute或/job:localhost/replica:0/task:0/gpu:0 Compute的行。某些流行和localhost/gpu行包含conv2D、relu或add、mul等层。然而，他们的时间表却大不相同。事实上，我不确定，我建议按照

generate\u chrome\u trace\u format

中的碎屑轨迹，发布你的发现--