37

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

Florin Andrei
  • 1,130
  • 1
  • 9
  • 13

4 Answers4

71

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.

Florin Andrei
  • 1,130
  • 1
  • 9
  • 13
16

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
tttzof351
  • 261
  • 2
  • 2
6

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.

momo
  • 163
  • 1
  • 5
0
CUDA_VISIBLE_DEVICES="-1" python3 your_file.py
khatchad
  • 101
  • 2