How to set a specific GPU in Tensorflow?

by coty_beier , in category: General Help , 7 months ago

How to set a specific GPU in Tensorflow?

Facebook Twitter LinkedIn Telegram Whatsapp

2 answers

Member

by tina , 7 months ago

@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.

Member

by tina , 2 months ago

@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.