让 TensorFlow 在 ARM Mac 上使用 GPU

Question

我已经根据these instructions 在 M1 (ARM) Mac 上安装了TensorFlow。一切正常。

但是，模型训练正在 CPU 上进行。如何将培训切换到GPU？

In: tensorflow.config.list_physical_devices()
Out: [PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]

在Apple's TensorFlow distribution 的文档中，我发现以下paragraph 有点令人困惑：

无需对现有的 TensorFlow 脚本进行任何更改即可将 ML Compute 用作 TensorFlow 和 TensorFlow Addons 的后端。有一个可选的mlcompute.set_mlc_device(device_name='any') API 用于 ML 计算设备选择。 device_name 的默认值为“any”，这意味着 ML Compute 将选择系统上可用的最佳设备，包括多 GPU 配置上的多个 GPU。其他可用选项是CPU 和GPU。请注意，在 Eager 模式下，ML Compute 将使用 CPU。例如，要选择 CPU 设备，您可以执行以下操作：

# Import mlcompute module to use the optional set_mlc_device API for device selection with ML Compute.
from tensorflow.python.compiler.mlcompute import mlcompute

# Select CPU device.
mlcompute.set_mlc_device(device_name='cpu') # Available options are 'cpu', 'gpu', and 'any'.

所以我尝试运行：

from tensorflow.python.compiler.mlcompute import mlcompute
mlcompute.set_mlc_device(device_name='gpu')

然后得到：

WARNING:tensorflow: Eager mode uses the CPU. Switching to the CPU.

此时我被卡住了。如何将 GPU 上的 keras 模型训练到我的 MacBook Air？

TensorFlow 版本：2.4.0-rc0

Answer 1

更新

tensorflow_macos tf 2.4 存储库已由所有者存档。对于tf 2.5，请参考here。

---

完全禁用急切执行可能没有用，但tf. functions。试试这个并检查你的 GPU 使用情况，警告信息可以是misleading。

import tensorflow as tf
tf.config.run_functions_eagerly(False)

---

Mac-optimized TensorFlow 的当前版本有几个尚未修复的问题 (TensorFlow 2.4rc0)。最终，Eager 模式是 TensorFlow 2.x 中的默认行为，并且在 TensorFlow-MacOS 中也没有改变。但与官方不同的是，这个优化版本强制使用 CPU 进行 Eager 模式。正如他们所说的here。

...在 Eager 模式下，ML Compute 将使用 CPU。

这就是为什么即使我们明确设置了device_name='gpu'，它仍然会切换回 CPU，因为 Eager 模式仍处于开启状态。

from tensorflow.python.compiler.mlcompute import mlcompute
mlcompute.set_mlc_device(device_name='gpu')

WARNING:tensorflow: Eager mode uses the CPU. Switching to the CPU.

Disabling the eager mode 可以让程序利用 GPU，但这不是一般行为，可能会导致这样的puzzling performance on both CPU/GPU。目前，最合适的方法是选择device_name='any'，ML Compute 将查询系统上的可用设备并选择最佳设备来训练网络。 p>

Answer 2

尝试关闭急切执行... 通过关注

import tensorflow as tf

tf.compat.v1.disable_eager_execution()

让我知道它是否有效。