33
$\begingroup$

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

$\endgroup$

3 Answers 3

56
$\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$
10
  • 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
  • $\begingroup$ Thanks you so much for this! $\endgroup$
    – AIM_BLB
    Jul 25, 2020 at 9:52
  • 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
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
4
$\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, privacy policy and cookie policy

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