限制TensorFlow的显存占用


发布于

|

分类

发现TensorFlow会把所有GPU全部占满,即使什么都不跑。当跑着一个TensorFlow的时候,其他程序如果申请显存都会报显存不足。

所以需要限制一下。

方法一: 每个GPU占用百分之X

sess = tf.Session(config=tf.ConfigProto(gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.333)))

如果有3块12G的卡,那么TF会在每个卡上划分4G出来。当某块卡剩余显存不够4G,直接报错。

方法二: 不预先划分

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

显存动态申请。


评论

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注