只为 Keras（TF 后端）脚本分配一个 GPU

Question

我有一台带有 2 个 GPU 的机器。

通常，一个用于生产（即使用已经训练的模型进行预测），而另一个用于训练和试验新模型。

当我使用 theano 时，通过指定如下标志，我可以在一个 GPU 上运行我的脚本没有问题

THEANO_FLAGS="device=cuda0" training_script.py THEANO_FLAGS="device=cuda1" prediction_script.py

有没有一种简单的方法可以在 Keras 中使用 Tensorflow 后端做同样的事情？默认行为似乎映射了一个会话的所有 GPU 的所有内存

（请注意，我并不真正关心每个脚本是否单独映射整个 GPU，即使它们可以使用更少的内存工作）

Answer 1

您可以轻松选择一个 GPU。只需填写 0 或 1CUDA_VISIBLE_DEVICES

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

此外，如果您想为上面选定的 gpu 指定 gpu 的一部分，请添加：

from keras import backend as K
import tensorflow as tf
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.4 #what portion of gpu to use
session = tf.Session(config=config)
K.set_session(session)