How to set a specific GPU in Tensorflow?

@coty_beier 

In TensorFlow, you can use the tf.config.experimental.set_visible_devices function to specify which GPU you want to use. Here's the general process:

1. Import TensorFlow:

1

import tensorflow as tf

1. Get the list of all available physical GPUs:

1

gpus = tf.config.list_physical_devices('GPU')

1. Choose the GPU you want to use and assign it to the environment variable CUDA_VISIBLE_DEVICES. Make sure to replace X with the index number of the GPU you want to use:

1

tf.config.experimental.set_visible_devices(gpus[X], 'GPU')

Here's an example of setting the first GPU as the desired device:

1
2
3
4

import tensorflow as tf

gpus = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_visible_devices(gpus[0], 'GPU')

This will make TensorFlow only use the specified GPU for training or running computations.

@coty_beier 

It's important to note that the above approach works for TensorFlow versions below 2.6. For TensorFlow 2.6 and later, the physical_device function should be used. Here's an example:

1
2
3
4

import tensorflow as tf

physical_devices = tf.config.list_physical_devices('GPU')
tf.config.set_visible_devices(physical_devices[1], 'GPU')

Replace [1] in physical_devices[1] with the index of the GPU you wish to use. Additionally, in TensorFlow 2.1 and onwards, you can set the GPU to be the one visible with the following method:

1
2

physical_devices = tf.config.list_physical_devices('GPU')
tf.config.set_visible_devices(physical_devices[1], 'GPU')

As always, it's recommended to refer to the official TensorFlow documentation for the most current and accurate information.