Back to Home

Accelerate tensorflow model training on Macs using tensorflow-metal

Tensorflow, by default, uses your machine’s CPU to train models

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’s very time consuming to train (either because the model is complex or because the dataset is large or both), training on your machine’s GPU may be faster.

If you’re on a Mac, you can configure tensorflow to use your machine’s GPU to train your model. Here’s how.

Install the tensorflow-metal PluggableDevice

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

Tensorflow pluggable devices is a plugin system in tensorflow that gives tensorflow users the ability to configure tensorflow to perform computations on a custom device (The Mac’s GPU in this case).

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

Built with Hugo & Notion. Source code is available at GitHub.