How to set a specific GPU in Tensorflow?

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

How to set a specific GPU in Tensorflow?

Facebook Twitter LinkedIn Telegram Whatsapp

1 answer

Member

by tina , 3 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.