Accelerate tensorflow model training on Macs using tensorflow-metal
May 13, 2024
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