Question

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

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

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

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

在Keras中使用Tensorflow后端是否有一种简单的方法来做到这一点？默认行为似乎会映射一个会话的所有GPU的所有内存

（请注意，即使它们可以使用更少的内存工作，我也不在乎每个脚本是否分别映射整个GPU）

Answer 1

您可以轻松选择一个GPU。只需在CUDA_VISIBLE_DEVICES

上填写0或1

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)

仅将一个GPU分配给Keras（TF后端）脚本

1 个答案: