python - 让 TensorFlow 在 ARM Mac 上使用 GPU

Question

我已安装 TensorFlow在 M1 ( ARM ) Mac 上根据 these instructions .一切正常。
但是，模型训练正在CPU 上进行。 .如何将培训切换到 GPU ?

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

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

It is not necessary to make any changes to your existing TensorFlow scripts to use ML Compute as a backend for TensorFlow and TensorFlow Addons. There is an optional mlcompute.set_mlc_device(device_name='any') API for ML Compute device selection. The default value for device_name is 'any', which means ML Compute will select the best available device on your system, including multiple GPUs on multi-GPU configurations. Other available options are CPU and GPU. Please note that in eager mode, ML Compute will use the CPU. For example, to choose the CPU device, you may do the following:

# 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.

在这一点上，我被卡住了。我如何训练 keras GPU 上的模型到我的 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 )。最终，急切模式是 TensorFlow 2.x 中的默认行为，而且在 TensorFlow-MacOS 中也没有变化.但与官方不同的是，这个优化版使用了 CPU 强制为急切模式。正如他们所说 here .

... in eager mode, ML Compute will use the 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' ，由此 机器学习计算 将查询系统上可用的设备并选择最佳设备来训练网络。