36
$\begingroup$

Using tensorflow-gpu 2.0.0rc0. I want to choose whether it uses the GPU or the CPU.

$\endgroup$

3 Answers 3

69
$\begingroup$

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.

$\endgroup$
11
  • 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$
    – Djib2011
    Sep 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 corresponding Python runtime was still consuming graphics memory and the GPU fans turned ON when I executed my code. $\endgroup$
    – hafiz031
    Nov 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$
    – Darqer
    Nov 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
15
$\begingroup$

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
$\endgroup$
3
  • 1
    $\begingroup$ RuntimeError: "Visible devices cannot be modified after being initialized" $\endgroup$
    – rjurney
    Oct 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
6
$\begingroup$

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.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.