tensorflow显存管理

Posted 2020-11-05 张乐乐章

在运行上面的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)