Using tensorflow-gpu 2.0.0rc0. I want to choose whether it uses the GPU or the CPU.
3 Answers
I've seen some suggestions elsewhere, but they are old and do not apply very well to newer TF versions. What worked for me was this:
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
When that variable is defined and equal to -1, TF uses the CPU even when a CUDA GPU is available.
-
2$\begingroup$ I actually find it more convenient to set the environmental variable from outside the script. Sometimes I might forget it in a script and when I import something from that script it automatically runs it and I have no GPU. $\endgroup$– Djib2011Sep 7, 2019 at 22:39
-
1$\begingroup$ At 2.1 (or possibly before) up to nightly, set that environment variable to an empty string to disable GPUs $\endgroup$ May 21, 2020 at 23:09
-
2$\begingroup$ TensorFlow still uses GPU even after adding this snippet. I spotted it by running
nvidia-smi
command from the terminal. The correspondingPython
runtime was still consuming graphics memory and the GPU fans turned ON when I executed my code. $\endgroup$– hafiz031Nov 20, 2020 at 22:21 -
1$\begingroup$ After a long fight with "Failed to get device attribute 13 for device 0" error at the computer with weak graphic card, it fixed the issue. Thank you. $\endgroup$– DarqerNov 30, 2020 at 14:24
-
2$\begingroup$ Make sure to run this before importing tensorflow. If you run this afterwards, it will use GPU even if this variable is set. $\endgroup$ Jul 11, 2022 at 14:48
For TF2:
try:
# Disable all GPUS
tf.config.set_visible_devices([], 'GPU')
visible_devices = tf.config.get_visible_devices()
for device in visible_devices:
assert device.device_type != 'GPU'
except:
# Invalid device or cannot modify virtual devices once initialized.
pass
-
1$\begingroup$ RuntimeError: "Visible devices cannot be modified after being initialized" $\endgroup$– rjurneyOct 29, 2020 at 19:25
-
$\begingroup$ @rjurney make sure you call this code before any tensorflow op $\endgroup$ Nov 18, 2020 at 14:02
-
$\begingroup$ This is the only way that looks to work with AMD / ROCm GPUs. Thank you! $\endgroup$ Nov 6, 2021 at 17:59
I find setting the variable outside the script easiest and something that always works.
export CUDA_VISIBLE_DEVICES=''
Run this on the command line before running your python script.