设置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’

上一篇:在IE环境下用javascript触发ABAP event


下一篇:Deep Learning Compiler 之自我理解