Tensorflow - 使用 CPU 版本 [重复]

Question

我在我的 anaconda 基础环境中安装了 tensorflow 包。我注意到从 2.1 版开始，它包含 CPU 和 GPU 包。 现在，在导入时，语法将保持不变， import tensorflow as tf

但是如果我只想使用 CPU 版本（当网络很小并且在 GPU 上运行可能会更慢）或 GPU 版本（其余的），我该如何指定呢？

Answer 1

在计算时选择合适的设备似乎更简单。可以像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)

要为您的案例获取正确的/CPU:0 字符串，请使用tf.config.experimental.list_physical_devices('CPU')。它将返回一个包含正确设备名称的列表。

在运行代码之前，您也可以在 ubuntu 终端中执行 export CUDA_VISIBLE_DEVICES=""。它将阻止使用任何 GPU 资源。见this answer。