Is there a way to change GPU / CPU on tensorflow during code execution?

Question

I have 2 tensorflow (1.15.4) models running sequentially. The output from the first model will be fed into the second model. Is there a way to run the first model using CPU and run the second one using GPU in one python script?

To simplify it, I tried a sample script like below

import os
import tensorflow as tf
os.environ["CUDA_VISIBLE_DEVICES"]="-1"
print(tf.test.is_gpu_available())
os.environ.pop("CUDA_VISIBLE_DEVICES")
os.environ["CUDA_VISIBLE_DEVICES"]="0"
print(tf.test.is_gpu_available())

However, both print results are False .

Is it possible to switch the device type in the middle of execution of a piece of codes?

Thank you!

---

Update:

According the comments, I tried the tf.device as follows

import tensorflow as tf

with tf.device("/CPU:0"):
    code1

with tf.device('GPU:0'):
    code2

However, both code1 and code2 are run on CPU. Is there anything I am missing?

Thank you

Answer 1

GPU cards are not installed properly. That's why you received False in both prints. If the GPU card is installed properly, then tf.test.is_gpu_available will return True.

Since you are using TF 1.15.4, please install Tensorflow GPU as shown below

!pip install tensorflow_gpu==1.15.4

Please refer to the sample working code to use both CPU and GPU as shown below

import tensorflow as tf
print(tf.__version__)
print(tf.test.is_gpu_available())


# Place tensors on the CPU
with tf.device('/CPU:0'):
  a1 = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
  a2 = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  a3 = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

# Run on the CPU
with tf.device('/CPU:0'):
  b = tf.add(a1, a3)

print(c)

# Run on the GPU
with tf.device('/GPU:0'):
  c = tf.matmul(a2, a3)

print(d)

Output:

1.15.4
True
Tensor("Add_1:0", shape=(3, 2), dtype=float32, device=/device:CPU:0)
Tensor("MatMul_1:0", shape=(2, 2), dtype=float32, device=/device:GPU:0)

For more details please refer here to install GPU and here for tested build configurations.