[TF 2.0] How to globally force CPU?

[jbgh2]

In TF 1.x it was possible to force CPU only by using:

config = tf.ConfigProto(device_count = {'GPU': 0})

However, ConfigProto doesn't exist in TF 2.0 and changing a OS environment variable seems very clunky.

What's the TF 2.0 way of doing this?

[pandrey-fr]

There might be a better way, but you can use tf.config.experimental.set_visible_devices([], 'GPU') to hide any GPU (they can still be listed using tf.config.experimental.list_physical_devices('GPU') and restored using the first function with different arguments, but will no longer appear when running tf.config.experimental.list_logical_devices('GPU') and should therefore not be used to place operations).

[ymodak]

This can help you set all operations on CPU;

import tensorflow as tf
print(tf.__version__)

# Set CPU as available physical device
my_devices = tf.config.experimental.list_physical_devices(device_type='CPU')
tf.config.experimental.set_visible_devices(devices= my_devices, device_type='CPU')

# To find out which devices your operations and tensors are assigned to
tf.debugging.set_log_device_placement(True)

# Create some tensors and perform an operation
a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
c = tf.matmul(a, b)

print(c)

Output

2.0.0-beta1
Executing op MatMul in device /job:localhost/replica:0/task:0/device:CPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

[jbgh2]

That works, thank you.

[ymodak]

Are you satisfied with the resolution of your issue?
Yes
No

[henrysky]

for anyone who is using tf 2.1, the above comment does not seems to work.

I use tf.config.set_visible_devices([], 'GPU') seems to do the job correctly

[pandrey-fr]

@henrysky this is indeed part of the breaking changes introduced in 2.1 (as documented in the release notes), which promoted some former experimental functions to the public scope; thank you for pointing the new correct syntax out :)

[JackXueIndiana]

tf.config.set_visible_devices([], 'GPU')
Works for tf 2.2

[gregoruar]

I managed to limit cpus by using psutils lib.
I provided this the beginning of the function

pid = psutil.Process(os.getpid())
pid.cpu_affinity([0, 1])

The later call of model.fit utilized exactly 2 cpus

[marcusobrien]

I also wanted to use TF 2.10 with CPU only as it defaults to using my GPU. However the above can not be used in Jupyter Labs due to this error message

"RuntimeError: Visible devices cannot be modified after being initialized"

So is there any other way to force tensorflow to use CPU to perform all calculations ?