python-3.x - 只分配一个GPU给Keras(TF后端)脚本

Question

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

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

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

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

有没有一种简单的方法可以在 Keras 中使用 Tensorflow 后端执行相同的操作？默认行为似乎映射一个 session 的所有 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)