只为 Keras（TF 后端）脚本分配一个 gpu答案

【问题标题】：Allocate only one gpu to Keras (TF backend) script只为 Keras（TF 后端）脚本分配一个 gpu
【发布时间】：2018-12-31 01:23:20
【问题描述】：

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

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

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

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

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

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

【问题讨论】：

标签： python-3.x tensorflow machine-learning keras gpu

【解决方案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)

【讨论】：