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()
config.gpu_options.per_process_gpu_memory_fraction = 0.4
set_session(tf.Session(config=config))

需要注意的是，虽然代码或配置层面设置了对显存占用百分比阈值，但在实际运行中如果达到了这个阈值，程序有需要的话还是会突破这个阈值。换而言之如果跑在一个大数据集上还是会用到更多的显存。以上的显存限制仅仅为了在跑小数据集时避免对显存的浪费而已。（悟乙己大佬提供）

3、指定GPU

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

此时的代码为选择了编号为1的GPU

4、tensorflow + CPU充分使用

num_cores = 4

config = tf.ConfigProto(intra_op_parallelism_threads=num_cores, inter_op_parallelism_threads=num_cores,
                        allow_soft_placement=True, device_count={'CPU': 4})
session = tf.Session(config=config)
K.set_session(session)

其中：

  1. device_count, 告诉tf Session使用CPU数量上限，如果你的CPU数量较多，可以适当加大这个值
  2. inter_op_parallelism_threads和intra_op_parallelism_threads告诉session操作的线程并行程度，如果值越小，线程的复用就越少，越可能使用较多的CPU核数。如果值为0，TF会自动选择一个合适的值。
  3. allow_soft_placement=True，
    有时候，不同的设备，它的cpu和gpu是不同的，如果将这个选项设置成True，那么当运行设备不满足要求时，会自动分配GPU或者CPU。

5、tf.keras使用多GPU

DistributionStrategy API是构建多设备/机器训练的简单方式，开发者只需要在现有模型上做少量的修改，就可以用它们进行分布式训练。另外，DistributionStrategy在设计时考虑了同时兼容动态图(eager)和静态图。
目前TensorFlow支持三种DistributionStrategy：

  1. MirroredStrategy
  2. CollectiveAllReduceStrategy
  3. ParameterServerStrategy

在tf.keras中直接使用DistributionStrategy

最新的TensorFlow Github中给出了在tf.keras中直接使用DistributionStrategy的例子。

用tf.keras构建一个单层网络：

inputs = tf.keras.layers.Input(shape=(1,))
predictions = tf.keras.layers.Dense(1)(inputs)
model = tf.keras.models.Model(inputs=inputs, outputs=predictions)

目前，使用DistributionStrategy需要使用tf.data.Dataset来作为数据输入：

features = tf.data.Dataset.from_tensors([1.]).repeat(10000).batch(10)
labels = tf.data.Dataset.from_tensors([1.]).repeat(10000).batch(10)
train_dataset = tf.data.Dataset.zip((features, labels))

这里我们为模型指定使用MirroredStrategy进行多GPU训练，代码非常简单：

distribution = tf.contrib.distribute.MirroredStrategy()
model.compile(loss='mean_squared_error',
              optimizer=tf.train.GradientDescentOptimizer(learning_rate=0.2),
              distribute=distribution)

使用常规的训练、评价和预测方法会自动在多GPU上进行：

model.fit(train_dataset, epochs=5, steps_per_epoch=10)
model.evaluate(eval_dataset)
model.predict(predict_dataset)

将tf.keras模型迁移到多GPU上运行只需要上面这些代码，它会自动切分输入、在每个设备（GPU）上复制层和变量、合并和更新梯度。