使用Tensorflow或Keras时对GPU内存限制

博客原文——使用Tensorflow或Keras时对GPU内存限制

跑Keras 或者 Tensorflow时默认占满所有GPU内存,这时如果想再开一个进程,或者别人想开一个进程都挤不上来,所以必须限制GPU内存

最好的资料还是官方文档

  • visible_device_list指定使用哪块显卡
  • per_process_gpu_memory_fraction分配到的内存占总内存量的比例
  • allow_growth根据运行时的需要来分配GPU内存,刚开始分配很少内存,随着Session开始运行并需要更多GPU内存时会自动扩展,但后续不会释放内存
  • per_process_gpu_memory_fractionallow_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

你可能感兴趣的:(使用Tensorflow或Keras时对GPU内存限制)