对于 tensorflow 2.x，如何在 CPU 和 GPU 版本之间切换?

Question

我已经将anaconda更新到最新版本。它同时安装了 tensorflow 2.2 和 tensorflow-gpu 2.2。但是当我import tensorflow时，张量流-gpu是默认使用的。有没有办法在它们之间切换？

本网站上提出的类似问题都是针对1.x版本的。我已经尝试了解决方案，但似乎它不适用于 2.x 版本。在 CPU/GPU 之间切换张量流

Answer 1

这应该可以解决您的问题：

with tf.device('/gpu:0'):
# YOUR def main() OR model.fit()

with tf.device('/cpu:0'):
# YOUR def main() OR model.fit()

这应该适用于没有会话的 TF2。

Answer 2

另一种方式是：

#use gpu 0
physical_devices = tf.config.list_physical_devices('GPU')
tf.config.set_visible_devices(physical_devices[0], 'GPU')
#use cpu
tf.config.set_visible_devices([], 'GPU')