Accelerate tensorflow model training on Macs using tensorflow-metal

Tensorflow, by default, trains models on your Mac’s CPU

python -c "import tensorflow as tf; print(tf.config.list_physical_devices())"

[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]

However, if you’re training a model that takes a long time to train (either because the model has many parameters or because the training dataset is large or both), training the model on your Mac’s GPU can be faster.

Here’s how you configure tensorflow to train your model on your Mac’s GPU.

Install the tensorflow-metal tensorflow pluggable device. Tensorflow pluggable devices is a plugin system in tensorflow that gives tensorflow users the ability to configure tensorflow to perform computations on any device (like the Mac’s GPU).

pip install tensorflow-metal 

Alternatively, if you’ve a requirements.txt file for your Python project, you can add tensorflow-metal package

# requirements.txt

tensorflow
tensorflow-metal

pip install -r requirements.txt

Once the tensorflow-metal package is installed, tensorflow will start using your Mac’s GPU to train models and training epochs should be faster

python -c "import tensorflow as tf; print(tf.config.list_physical_devices())"

[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'), PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]

If you’d like to learn more, see Apple’s documentation on tensorflow-metal and tensorflow’s documentation on pluggable devices

https://developer.apple.com/metal/tensorflow-plugin/ https://blog.tensorflow.org/2021/06/pluggabledevice-device-plugins-for-TensorFlow.html