Heven

本文转载自知乎

Keras使用显卡时是默认调用所有的GPU，并且占满所有显存的，所以就很有必要搞清楚Keras如何指定GPU和如何限制显存的使用比例了。

📒指定某块GPU

指定GPU很简单，在载入keras和tensorflow之前，设置CUDA计算使用的GPU序号即可。这其实是CUDA本身的参数，对所有深度学习框架都是适用的。需要注意的是最好写在improt keres 和 import tensorflow之前，不然可能出错。

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

此代码选择的是编号为“1”的显卡。显卡的编号是从“0”开始的，若只有一块显卡则编号为“0”。

📒指定多块GPU

指定多块GPU的方式和前文完全一致，只需要多写几个编号即可：

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

此时选择的是编号为“0”和编号为“2”的显卡

📒控制GPU显存使用比例

Keras是默认占满GPU显存的，我们通过重设backend的gpu_memory_fraction来进行调节，0.3表示占用30%的显存：

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))

不过虽然配置了GPU显存的占用比例，实际运行中若不够用的话还是获取更多的显存。比方说如果运行了3个设置为30%显存的应用，实际上是可能超过100%造成显存不足的。

📒指定GPU+控制显存使用

将指定GPU与控制显存使用比例合并操作即可：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
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))

该代码实现的是指定编号为“1”的GPU，并设置占用显存的比例为30%。

📒显存的按需分配（动态增长）

如果并不清楚自己的应用分配多少的显存比例合适，可以使用按需分配的方式，也就是动态增长allow_growth：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
set_session(tf.Session(config=config))

该代码实现的是指定编号为“1”的GPU，并设置占用显存的方式为按需增长。