keras控制显存(GPU)

Tensorflow与Keras在使用GPU进行训练的时候默认会申请全部的显存。
这造成了计算资源的一种巨大浪费。
下面介绍几种常见的处理方案：

1、GPU动态增长

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


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

config = tf.ConfigProto()
config.gpu_options.allow_growth=True   #不全部占满显存, 按需分配
sess = tf.Session(config=config)

KTF.set_session(sess)
os.environ指的是占用的GPU编号；allow_growth为动态申请显存占用。

2、固定显存的GPU

当许多人使用单张卡的时候，需要限制一下使用的显存：

import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto