Allocate only one gpu to Keras (TF backend) script

Question

I have a machine with 2 GPUs.

Quite often, one is used in production (i.e doing predictions with the already trained model), while the other is used for training and experimenting new models.

While I was using theano, I had no problem running my scripts on only one GPU by specifying a flag as follow

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

Is there a simple way to do the same in Keras with a Tensorflow backend ? Default behavior seem to map all the memory of all the GPUs for one session

(Please note that I don't really care if each script maps a whole GPU separately, even if they could work using less memory)

Answer 1

You can easily choose one gpu. Just fill 0 or 1 on CUDA_VISIBLE_DEVICES

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

Furthermore if you want to spesify a portion of gpu for the selected gpu above, add:

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)