我正在使用带有 Tensorflow 后端的 Keras 并查看nvidia-smi
不足以了解当前网络架构需要多少内存,因为 Tensorflow 似乎只是分配所有可用内存。
那么问题是如何找出真实的GPU内存使用情况?
这可以使用时间轴来完成,它可以为您提供有关内存记录的完整跟踪。类似于下面的代码:
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_options来完成。就像下面的代码:
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))
检查以下有关的文档Timeline https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/client/timeline.py object
当您在后端使用 TensorFlow 时,您可以使用tfprof https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/tfprof/README.md分析工具
本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系:hwhale#tublm.com(使用前将#替换为@)