Python Keras：实际使用的GPU内存量

我在Tensorflow后端使用Keras，查看

nvidia smi

不足以了解当前网络架构需要多少内存，因为Tensorflow似乎只分配所有可用内存

---

所以问题是如何找出真正的GPU内存使用情况

可以使用Timeline完成，它可以为您提供有关内存日志记录的完整跟踪。与下面的代码类似：

from keras import backend as K
from tensorflow.python.client import timeline
import tensorflow as tf


with K.get_session()  as s:
    run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
    run_metadata = tf.RunMetadata()
     
    # your fitting code and s run with run_options 

    to = timeline.Timeline(run_metadata.step_stats)
    trace = to.generate_chrome_trace_format()
    with open('full_trace.json', 'w') as out:
            out.write(trace)

如果你想限制gpu内存的使用，也可以通过gpu选项来实现。如以下代码所示：

import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.2
set_session(tf.Session(config=config))

检查以下有关该对象的文档

---

当您在后端使用TensorFlow时，您可以使用分析工具

在告诉TensorFlow不要保留GPU的所有内存，而是根据需要增加此保留后，您仍然可以使用

nvidia smi

：

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
keras.backend.tensorflow_backend.set_session(tf.Session(config=config))

---

您试用过“model.summary（）”吗？它应该给出一些关于模型内存使用的想法。@orabis是的，但它只是权重，如果我们训练模型Tensorflow也分配itermediate Blob和gradients Blob+一些开销，我不知道如何精确计算内存使用。问题是如果你运行

fit

，然后是

会话。运行
，当插入指令的
run
调用开始时，
fit
分配的大部分内存将被释放。这里有一个相关的问题：缺少的是让Keras使用自定义
run\u选项
，或者向Keras模型添加自定义操作的方法（比如
MaxBytesInUse
op-like）