Tensorflow与Keras自适应使用显存

Tensorflow支持基于cuda内核与cudnn的GPU加速，Keras出现较晚，为Tensorflow的高层框架，由于Keras使用的方便性与很好的延展性，之后更是作为Tensorflow的官方指定第三方支持开源框架。但两者在使用GPU时都有一个特点，就是默认为全占满模式。在训练的情况下，特别是分步训练时会导致显存溢出，导致程序崩溃。可以使用自适应配置来调整显存的使用情况。

Tensorflow

指定CPU

import os
os.environ["CUDA_VISIBLE_DEVICES"] = " "

# 如果想要指定CPU核数，使用如下方法 
cpu_num = 10#指定使用的CPU个数
config = tf.ConfigProto(device_count={"CPU": cpu_num},
            inter_op_parallelism_threads = cpu_num,
            intra_op_parallelism_threads = cpu_num,
            log_device_placement=True)
# 开始训练
with tf.Session(config=config) as sess:
    #以下编写自己的代码

指定显卡

代码中加入

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# os.environ["CUDA_VISIBLE_DEVICES"]='6, 7'
print('CUDA Device: ' + os.environ["CUDA_VISIBLE_DEVICES"])

或者在运行代码前，在终端

export CUDA_VISIBLE_DEVICES=0

为显存分配使用比例

在建立tf.Session加入设置数据（显存使用比例为1/3），但有时你虽然设置了使用上限，在程序需要更高显存时还是会越过该限制

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)  
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

自适应分配

会自适应分配显存，不会将显存全部分配导致资源浪费

config = tf.ConfigProto()  
config.gpu_options.allow_growth=True  
sess = tf.Session(config=config)

Keras

与tensorflow大差不差，就是将tf.Session配置转置Keras配置

指定显卡

代码中加入

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

或者在运行代码前，在终端

export CUDA_VISIBLE_DEVICES=0

为显存分配使用比例

import tensorflow as tf
import keras.backend.tensorflow_backend as KTF

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.333
session = tf.Session(config=config)
KTF.set_session(session)

自适应分配

import keras.backend.tensorflow_backend as KTF

config = tf.ConfigProto()  
config.gpu_options.allow_growth=True   
session = tf.Session(config=config)
KTF.set_session(session)

设置fit_generator

如有设置fit_generator，将多线程关闭

#可将
use_multiprocessing=True
#改为
use_multiprocessing=False

参考

TensorFlow指定特定GPU以及占用显存的比例
TensorFlow之CPU、GPU设置