tensorflow显存管理

在运行上面的blog的Tensorflow小程序的时候程序我们会遇到一个问题,当然这个问题不影响我们实际的结果计算,但是会给同样使用这台计算机的人带来麻烦,程序会自动调用所有能调用到的资源,并且全占满,在自己的PC上没问题,但是在服务器上,问题就很大,因为一旦你运行程序,占满显存别人就不能再用了,解决方法是我们通常是给程序运行指定一块GPU,然后限制其使用比例:

1 import tensorflow as tf  
2 import os  
3 os.environ["CUDA_VISIBLE_DEVICES"] = '0'   #指定第一块GPU可用  
4 config = tf.ConfigProto()  
5 config.gpu_options.per_process_gpu_memory_fraction = 0.5  # 程序最多只能占用指定gpu50%的显存  
6 sess = tf.Session(config = config) 

但是我们通常在程序开始之前并不知道需要多大的显存,程序回去申请GPU的显存的50%,比如一个8G的内存,被占用了2G,那么程序会申请4G的显存(因为有足够的剩余显存),如果此时显存被占用7G,那么程序会申请剩下的所有的1G的显存。也许你的程序根本用不着这么多显存,200M也许就够了,这时候如果程序能按照需求去申请就好了,幸运的是,这样的方法是存在的:

只需要加上一条语句就好,完整语句为:

1 import tensorflow as tf  
2 import os  
3 os.environ["CUDA_VISIBLE_DEVICES"] = '0'   #指定第一块GPU可用  
4 config = tf.ConfigProto()  
5 config.gpu_options.per_process_gpu_memory_fraction = 0.5  # 程序最多只能占用指定gpu50%的显存  
6 config.gpu_options.allow_growth = True      #程序按需申请内存  
7 sess = tf.Session(config = config)  

猜你喜欢

转载自www.cnblogs.com/zle1992/p/9047888.html