有没有办法确定TensorFlow使用多少GPU内存？

Question

Tensorflow倾向于在其GPU上预分配整个可用内存.对于调试,有没有办法说明实际使用了多少内存？

Answer 1

(1)Timeline对记录内存分配的支持有限.以下是其用法示例:

    run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
    run_metadata = tf.RunMetadata()
    summary, _ = sess.run([merged, train_step],
                          feed_dict=feed_dict(True),
                          options=run_options,
                          run_metadata=run_metadata)
    train_writer.add_run_metadata(run_metadata, 'step%03d' % i)
    train_writer.add_summary(summary, i)
    print('Adding run metadata for', i)
    tl = timeline.Timeline(run_metadata.step_stats)
    print(tl.generate_chrome_trace_format(show_memory=True))
    trace_file = tf.gfile.Open(name='timeline', mode='w')
    trace_file.write(tl.generate_chrome_trace_format(show_memory=True))

您可以尝试使用MNIST示例尝试此代码(带有摘要的mnist)

这将生成一个名为timeline的跟踪文件,您可以使用chrome:// tracing打开该文件.请注意,这仅提供近似的GPU内存使用情况统计信息.它基本上模拟了GPU执行,但无法访问完整的图元数据.它也不知道已经为GPU分配了多少变量.

(2)对于非常粗略的GPU内存使用量,nvidia-smi将显示运行命令时的总设备内存使用量.

nvprof可以显示片上共享内存使用情况和CUDA内核级别的寄存器使用情况,但不显示全局/设备内存使用情况.

这是一个示例命令:nvprof --print-gpu-trace matrixMul

而更多的细节在这里: http://docs.nvidia.com/cuda/profiler-users-guide/#abstract

Answer 2

有一些代码tensorflow.contrib.memory_stats可以帮助解决这个问题：

from tensorflow.contrib.memory_stats.python.ops.memory_stats_ops import BytesInUse
with tf.device('/device:GPU:0'):  # Replace with device you are interested in
  bytes_in_use = BytesInUse()
with tf.Session() as sess:
  print(sess.run(bytes_in_use))

Answer 3

这是一个对我来说效果很好的实用解决方案：

使用 TF 会话配置禁用 GPU 内存预分配：

config = tf.ConfigProto()  
config.gpu_options.allow_growth=True  
sess = tf.Session(config=config)  

运行nvidia-smi -l（或其他一些实用程序）来监控 GPU 内存消耗。

使用调试器逐步执行您的代码，直到您看到意外的 GPU 内存消耗。