TensorFlow：如何记录 GPU 内存（VRAM）利用率？

Question

TensorFlow 总是（预）分配我显卡上的所有空闲内存 (VRAM)，这没关系，因为我希望我的模拟在我的工作站上尽可能快地运行。

但是，我想记录 TensorFlow 实际使用了多少内存（总计）。此外，如果我还可以记录单个张量使用的内存量，那就太好了。

此信息对于测量和比较不同 ML/AI 架构所需的内存大小很重要。

有什么建议吗？

Answer 1

更新，可以使用TensorFlow ops查询分配器：

# maximum across all sessions and .run calls so far
sess.run(tf.contrib.memory_stats.MaxBytesInUse())
# current usage
sess.run(tf.contrib.memory_stats.BytesInUse())

您还可以通过查看 RunMetadata 获得有关 session.run 调用的详细信息，包括在 run 调用期间分配的所有内存。 IE 是这样的

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

这是一个端到端的示例——获取列向量、行向量并将它们相加以获得加法矩阵：

import tensorflow as tf

no_opt = tf.OptimizerOptions(opt_level=tf.OptimizerOptions.L0,
                             do_common_subexpression_elimination=False,
                             do_function_inlining=False,
                             do_constant_folding=False)
config = tf.ConfigProto(graph_options=tf.GraphOptions(optimizer_options=no_opt),
                        log_device_placement=True, allow_soft_placement=False,
                        device_count={"CPU": 3},
                        inter_op_parallelism_threads=3,
                        intra_op_parallelism_threads=1)
sess = tf.Session(config=config)

with tf.device("cpu:0"):
    a = tf.ones((13, 1))
with tf.device("cpu:1"):
    b = tf.ones((1, 13))
with tf.device("cpu:2"):
    c = a+b

sess = tf.Session(config=config)
run_metadata = tf.RunMetadata()
sess.run(c, options=tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE, output_partition_graphs=True), run_metadata=run_metadata)
with open("/tmp/run2.txt", "w") as out:
  out.write(str(run_metadata))

如果您打开 run.txt，您会看到这样的消息：

  node_name: "ones"

      allocation_description {
        requested_bytes: 52
        allocator_name: "cpu"
        ptr: 4322108320
      }
  ....

  node_name: "ones_1"

      allocation_description {
        requested_bytes: 52
        allocator_name: "cpu"
        ptr: 4322092992
      }
  ...
  node_name: "add"
      allocation_description {
        requested_bytes: 676
        allocator_name: "cpu"
        ptr: 4492163840

所以这里可以看到a和b各分配了52字节（13*4），结果分配了676字节。

Answer 2

Yaroslav Bulatov 的回答是 TF1 的最佳解决方案。

但是，对于 TF2，contrib 包不存在。最好的方法是使用 tf 的分析器 -- https://www.tensorflow.org/guide/profiler#memory_profile_tool

它将绘制这样的内存利用率图。

TensorFlow：如何记录 GPU 内存（VRAM）利用率？

TensorFlow: how to log GPU memory (VRAM) utilization?

python

vram

tensorflow