How to set specific gpu in tensorflow?

Question

I want to specify the gpu to run my process. And I set it as follows:

import tensorflow as tf
with tf.device(\'/gpu:0\'):
    a = tf.constant(3.0)
with tf.Se         


        

Answer 1

You can modify the GPU options settings by adding at the begining of your python script:

gpu_options = tf.GPUOptions(visible_device_list="0")
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

"0" is here the name of the GPU you want to use. You can have the list of the GPU available by typing the command nvidia-smi in the terminal prompt.

---

With Keras, these 2 functions allow the selection of CPU or GPU and in the case of GPU the fraction of memory that will be used.

import os
from keras.backend.tensorflow_backend import set_session
import tensorflow as tf



def set_cpu_option():
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"  # see issue #152
    os.environ["CUDA_VISIBLE_DEVICES"] = ""
    os.environ["CUDA_VISIBLE_DEVICES"] = ""


def set_gpu_option(which_gpu, fraction_memory):
    config = tf.ConfigProto()
    config.gpu_options.per_process_gpu_memory_fraction = fraction_memory
    config.gpu_options.visible_device_list = which_gpu
    set_session(tf.Session(config=config))
    return

set_gpu_option("0", 0.9)
# or 
set_cpu_option()