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’