Running multiple graphs in different device modes in TensorFlow

Question

I am operating on multiple graphs simultanously. I would for example, like one graph to use CPU and the other to use the GPU. How can I achieve this ?

Current approach and its problems

When I use a tf.Session() with a tf.ConfigProto as follows, it does not work and still uses a GPU.

config = tf.ConfigProto(
            device_count = {'GPU': 0}
        )

I have to use the environment variable CUDA_VISIBLE_DEVICES to disable the use of a GPU. I later on use os.unsetenv() to remove this variable after my work.

These solutions are not useful to me because for one graph I want GPUs to be used and for the other I don't want the GPUs to be used. Setting os.environ() will affect both the graphs.

How can I achieve my purpose ?

Answer 1

The config created by tf.ConfigProto() can set the visible devices for the sessions. For exmaple:

config_cpu = tf.ConfigProto()
config_cpu.gpu_options.visible_device_list=''
sess_cpu = tf.Session(config=config_cpu)

config_gpu = tf.ConfigProto()
config_gpu.gpu_options.visible_device_list='0'
sess_gpu = tf.Session(config=config_gpu)

Then the graph in the session, sess_cpu, should run on CPU only and the graph in the session, sess_gpu, should run on GPU 0 only. To prevent tensorflow occupying the whole GPU memory, you can set config_gpu.gpu_options.allow_growth=True. Similar configurations can be adopted for your customized needs. You can take a look at the tf.ConfigProto if you want to use other configurations.