TensorFlow限制显存使用量

显式地设置显存使用比例

Tensorflow默认会完全占用满指定GPU的显存, 而线上部署后, 只需要进行推测, 是不需要占用完整块GPU显存的. 需要在代码中控制显存.

设置静态值

使用per_process_gpu_memory_fraction参数, 显式的指定每块GPU使用的内存比例, 设置的范围为0~1, 代表占用上限为0~100%.

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

keras需要使用set_session函数设置Session.

按需分配

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

使用allow_growth设置显存按需分配, 动态增长. 但注意一旦当前服务获取到了一定数量的显存, 后面是不会再还给系统的, 即使因为某些原因显存使用降低了. 总结来说, 这是一个不可逆的过程.

参考资料

• 限制TensorFlow显存使用量