博客原文——使用Tensorflow或Keras时对GPU内存限制
跑Keras 或者 Tensorflow时默认占满所有GPU内存，这时如果想再开一个进程，或者别人想开一个进程都挤不上来，所以必须限制GPU内存
最好的资料还是官方文档
-
visible_device_list指定使用哪块显卡 -
per_process_gpu_memory_fraction分配到的内存占总内存量的比例 -
allow_growth根据运行时的需要来分配GPU内存，刚开始分配很少内存，随着Session开始运行并需要更多GPU内存时会自动扩展，但后续不会释放内存 -
per_process_gpu_memory_fraction与allow_growth两种方法二选一即可。一般用allow_growth即可，如果你能准确预估你的程序需要多大显存，推荐per_process_gpu_memory_fraction
# 将下面代码加在.py文件开头即可
import tensorflow as tf
config = tf.ConfigProto()
config.gpu_options.visible_device_list = '0'
# config.gpu_options.per_process_gpu_memory_fraction = 0.2
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
# from keras import backend as K
# K.set_session(sess)
个人习惯而言，每个项目有个配置文件config.py，里面写好基本配置函数，如GPU配置，路径配置等。
set_gpu()函数，自动识别是否有显卡，有显卡则指定空余内存多或者GPU使用低的显卡，并设置分配比例
代码针对两个显卡的，懒得改了，多显卡类似啦~
def set_gpu(ratio=0, target='memory'):
"""
配置GPU，0表示自适应，(0, 1]表示显存占比
:param ratio:
:param target: 选择显存大的卡还是GPU利用率低的卡
:return:
"""
command1 = "nvidia-smi -q -d Memory | grep -A4 GPU | grep Free | awk '{print $3}'"
command2 = "nvidia-smi -q | grep Gpu | awk '{print $3}'"
memory = list(map(int, os.popen(command1).readlines()))
gpu = list(map(int, os.popen(command2).readlines()))
if memory and gpu: # 如果没有显卡，memory，gpu均为[]
if target == 'memory':
num = (1, 0)[memory[0] > memory[1]]
else:
num = (0, 1)[gpu[0] > gpu[1]]
print('>>> Free Memory : GPU0 %6d MiB | GPU1 %6d MiB' % (memory[0], memory[1]))
print('>>> Volatile GPU-Util : GPU0 %6d %% | GPU1 %6d %% ' % (gpu[0], gpu[1]))
print('>>> Using GPU%d' % num)
import tensorflow as tf
config = tf.ConfigProto()
config.gpu_options.visible_device_list = str(num) # 选择GPU
if ratio == 0:
config.gpu_options.allow_growth = True
else:
config.gpu_options.per_process_gpu_memory_fraction = ratio
sess = tf.Session(config=config)
from keras import backend as K
K.set_session(sess)
else:
print('>>> Could not find the GPU')
指定显卡的另外两种方法
- 配进环境变量
export CUDA_VISIBLE_DEVICES=0
- 使用os
import os
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
查看GPU使用情况
-
watch -n 1 nvidia-smi每秒刷新 nvidia-smi
