How to specify number of GPUs in Python interface?

Question

When I start a tensorflow Session it immediately allocates 100% of the memory on both GPUs in my dekstop (devices 0 and 2).

How can I use the Python API to restrict it to just 1 GPU?

Answer 1

A crude way to restrict it to use just GPU #0, for example, is to define this:

export CUDA_VISIBLE_DEVICES=0

Answer 2

If you want some operations to be done only with a particular device, you can do the following:

with tf.device('/gpu:0'):
   ...
   # all your operations

Where the devices are specified in the following way: "/cpu:0", "/gpu:0", "/gpu:1". I suggest you to read more details in the topic using GPU devices.

If you want to see some tutorials, take a look at cifar10 tutorial.