Make TensorFlow use the GPU on an ARM Mac

Question

I have installed TensorFlow on an M1 (ARM) Mac according to these instructions. Everything works fine.

However, model training is happening on the CPU. How do I switch training to the GPU?

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

In the documentation of Apple's TensorFlow distribution I found the following slightly confusing 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'.

So I try to run:

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

and get:

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

At this point I am stuck. How can I train keras models on the GPU to my MacBook Air?

TensorFlow version: 2.4.0-rc0

Answer 1

Update

The tensorflow_macos tf 2.4 repository has been archived by the owner. For tf 2.5, refer to here.

---

It's probably not useful to disable the eager execution fully but to tf. functions. Try this and check your GPU usages, the warning message can be misleading.

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

---

The current release of Mac-optimized TensorFlow has several issues that yet not fixed (TensorFlow 2.4rc0). Eventually, the eager mode is the default behavior in TensorFlow 2.x, and that is also unchanged in the TensorFlow-MacOS. But unlike the official, this optimized version uses CPU forcibly for eager mode. As they stated here.

... in eager mode, ML Compute will use the CPU.

That's why even we set explicitly the device_name='gpu', it switches back to CPU as the eager mode is still on.

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 may work for the program to utilize the GPU, but it's not a general behavior and can lead to such puzzling performance on both CPU/GPU. For now, the most appropriate approach can be to choose device_name='any', by that the ML Compute will query the available devices on the system and selects the best device(s) for training the network.

Answer 2

Try with turning off the eager execution... via following

import tensorflow as tf

tf.compat.v1.disable_eager_execution()

Let me know if it works.