设置GPU及显存大小

20210128 -

引言

之前搜索过设置GPU和显存大小的方式，但是升级了新的版本的keras以及tensorflow，导致之前的代码失效了，这里记录一下。

本质上，就是版本更换的原因，很多api可能被取消，或者改了别的。

原始代码

import os
import tensorflow as tf
from keras.backend.tensorflow_backend import set_session

def set_up_gpu_environment(gpu_devices_num = 2, memory_fraction = 1.0):
    available_gpu_num = ['0','1','2']
    if str(gpu_devices_num) not in available_gpu_num:
        print("The gpu deveices is not right")
    #设置GPU环境变量
    os.environ[ 'CUDA_VISIBLE_DEVICES'] = str(gpu_devices_num)
    #调整显存占用比例
    config = tf.ConfigProto()
    config.gpu_options.per_process_gpu_memory_fraction = memory_fraction
    set_session(tf.Session(config=config))

这部分代码我已经不知道是什么版本的tensorflow和keras了，上述代码主要由两部分，一部分通过设置环境变量来选择显卡，其中我使用的机器上有3个显卡，使用的时候，如果是全部，那么keras会占用所有的显卡，虽然在其他两块不工作的显卡上占用的显存不多，但还是通过环境变量来控制一下；另一部分工作是调整显存大小。

新版代码

新版的GPU环境按照之前的文章《Centos8安装英伟达显卡驱动并通过docker部署深度学习环境》，版本信息如下：

Keras 2.4.3
tensorflow-gpu 2.3.1

使用上述代码的时候会各种报错，例如第一条

ModuleNotFoundError: No module named 'keras.backend.tensorflow_backend'; 'keras.backend' is not a package

这条命令通过查询[1]，修改如下：

from tensorflow.python.keras.backend import set_session

但是这里通过之后，由于tensorflow的问题也会出错，之前的时候这部分问题也遇到过，所以就比较轻松的解决了，这次参考了文章[2]，那么整体上修改之后的代码如下：

import os
import tensorflow as tf
from tensorflow.python.keras.backend import set_session

def set_up_gpu_environment(gpu_devices_num = 2, memory_fraction = 1.0):
    available_gpu_num = ['0','1','2']
    if str(gpu_devices_num) not in available_gpu_num:
        print("The gpu deveices is not right")
    #设置GPU环境变量
    os.environ[ 'CUDA_VISIBLE_DEVICES'] = str(gpu_devices_num)
    #调整显存占用比例
    config = tf.compat.v1.ConfigProto()
    config.gpu_options.per_process_gpu_memory_fraction = memory_fraction
    set_session(tf.compat.v1.Session(config=config))

参考

[1]Import error: keras.backend.tensorflow_backend #197
[2]AttributeError: module ‘tensorflow’ has no attribute ‘ConfigProto’