使用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_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