1

I installed the tensorflow package in my anaconda base environment. I noticed that from version 2.1, it contains both CPU and GPU packages. Now, while importing, the syntax will stay the same, import tensorflow as tf

But if I want to use only the CPU version (when the network is small and running on GPU might be slower) or the GPU version (for the rest), how do I specify that?

malibu
  • 77
  • 3
  • 7

1 Answers1

1

It seems simpler to just select appropriate device when doing computation. It can be done like this:


# Place tensors on the CPU
with tf.device('/CPU:0'):
  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)

in order to get the correct /CPU:0 string for your case, use tf.config.experimental.list_physical_devices('CPU'). It will return a list with correct device names.

You could also do export CUDA_VISIBLE_DEVICES="" in ubuntu terminal, before running your code. It will prevent using any GPU resource. See this answer.

tornikeo
  • 915
  • 5
  • 20