tensorflow/keras指定显卡及限制GPU用量

作者： zydarChen
原文链接： https://www.zydarchen.top/20181016/22_tensorflow_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使用低的显卡，并设置分配比例

原文中的代码是针对2张显卡的，我这里改成了针对多张显卡的

def gpu_options(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 = memory.index(max(memory))
        else:
            num = gpu.index(min(gpu))
        print('>>> Using GPU%d' % num)
        print('>>> Free Memory       : %6d MiB' % (memory[num]))
        print('>>> Volatile GPU-Util : %6d %% ' % (gpu[num]))

        config = tf.compat.v1.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.compat.v1.Session(config=config)

        tf.compat.v1.keras.backend.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每秒刷新