keras是如何指定顯卡且限制顯存用量（GPU/CPU使用）

轉自（修訂）：素質雲博客 https://blog.csdn.net/sinat_26917383/article/details/75633754

參考：https://blog.csdn.net/zziahgf/article/details/80226129

keras在使用GPU的時候有個特點，就是默認全部佔滿顯存（多個顯卡）。
若單核GPU也無所謂，若是服務器GPU較多，性能較好，全部佔滿就太浪費了。
於是乎有以下五種情況：

1、指定GPU
2、使用固定顯存的GPU
3、指定GPU + 固定顯存
4 、GPU動態增長（按需分配）
5、 CPU充分佔用
一、固定顯存的GPU
本節來源於：深度學習theano/tensorflow多顯卡多人使用問題集（參見：Limit the resource usage for tensorflow backend · Issue #1538 · fchollet/keras · GitHub）
在使用keras時候會出現總是佔滿GPU顯存的情況，可以通過重設backend的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.3
set_session(tf.Session(config=config))

需要注意的是，雖然代碼或配置層面設置了對顯存佔用百分比閾值，但在實際運行中如果達到了這個閾值，程序有需要的話還是會突破這個閾值。換而言之如果跑在一個大數據集上還是會用到更多的顯存。以上的顯存限制僅僅爲了在跑小數據集時避免對顯存的浪費而已。（2017年2月20日補充）

二、指定GPU

注意指定GPU要在import tensorflw/keras 之前！否則可能無效。

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

此時的代碼爲選擇了編號爲2 的GPU

注意，在代碼中指定設備時，重新從0開始計，而不是從8開始。
來源：Tensorflow 學習筆記（七） ———— 多GPU操作

三、指定GPU + 固定顯存
上述兩個連一起用就行：

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

import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.3
set_session(tf.Session(config=config))

那麼在命令行，可以使用：https://github.com/tensorflow/nmt/issues/60

CUDA_VISIBLE_DEVICES=0 python -m nmt.nmt 

四、GPU動態增長（按需分配）

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1" # 注意在import keras/tensorflow之前

import keras.backend.tensorflow_backend as KTF
import tensorflow as tf
config = tf.ConfigProto()
config.gpu_options.allow_growth=True   #不全部佔滿顯存, 按需分配
sess = tf.Session(config=config)

KTF.set_session(sess)

os.environ指的是佔用的GPU編號；allow_growth爲動態申請顯存佔用。

五、tensorflow + CPU充分使用
來自博客：TensorFlow如何充分使用所有CPU核數，提高TensorFlow的CPU使用率，以及Intel的MKL加速

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)

其中：

device_count, 告訴tf Session使用CPU數量上限，如果你的CPU數量較多，可以適當加大這個值
inter_op_parallelism_threads和intra_op_parallelism_threads告訴session操作的線程並行程度，如果值越小，線程的複用就越少，越可能使用較多的CPU核數。如果值爲0，TF會自動選擇一個合適的值。
allow_soft_placement=True， 有時候，不同的設備，它的cpu和gpu是不同的，如果將這個選項設置成True，那麼當運行設備不滿足要求時，會自動分配GPU或者CPU。
六 tf.keras使用多GPU
DistributionStrategy API是構建多設備/機器訓練的簡單方式，開發者只需要在現有模型上做少量的修改，就可以用它們進行分佈式訓練。另外，DistributionStrategy在設計時考慮了同時兼容動態圖(eager)和靜態圖。
參考：TensorFlow 1.11.0發佈，一鍵多GPU（訓練、預測和評價tf.keras模型）

目前TensorFlow支持三種DistributionStrategy：

MirroredStrategy
CollectiveAllReduceStrategy
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）上覆制層和變量、合併和更新梯度。