How to train deep learning models on Apple Silicon GPU

Training deep learning models on GPUs is faster than training on the CPU, especially when training large image or language models. You can train deep learning models with GPUs on Apple Silicon devices.

To train deep learning models with GPU on Apple Silicon devices, you need to install the TensorFlow or PyTorch versions that are compatible with Apple Silicon devices. The next time you train a deep learning model, the training will be on the GPU by default.

Training deep learning models on Apple Silicon

In Apple Silicon rendering of advanced graphics in GPUs is done by Metal. The Metal framework gives apps direct access to GPU on Mac computers. You can leverage this framework to:

• Render games in complex environments.
• Process large datasets.
• Video editing.
• Training deep learning models on the GPU.

TensorFlow introduced the PluggableDevice architecture to enable the running of TensorFlow on new devices without changing TensorFlow code. The PluggableDevice architecture allows integration with new devices without changing TensorFlow core code.

TensorFlow

Training TensorFlow deep learning networks on Apple Silicon is done through the PluggableDevice architecture. We, therefore, install the compatible TensorFlow version and the Tensorflow-metal PluggableDevice.

Install Tensorflow-metal PluggableDevice

Let's walk through the installation instructions. Start by installing Conda and creating a virtual environment where TensorFlow will be installed. For this to work, ensure that you are using Python 3.8 or 3.9.

chmod +x ~/Downloads/Miniforge3-MacOSX-arm64.sh
sh ~/Downloads/Miniforge3-MacOSX-arm64.sh
source ~/miniforge3/bin/activate

Next, install the TensorFlow dependencies in this environment:

conda install -c apple tensorflow-deps

Install TensorFlow:

python -m pip install tensorflow-macos

Finally, install the tensorflow-metal plugin.

python -m pip install tensorflow-metal

You are now ready to train deep learning models on Apple Silicon with TensorFlow.

Train TensorFlow model on Apple Silicon GPU

After performing the above instructions, deep learning models will be trained by default on the GPU.

You can check the activity monitor to confirm that Python is using the GPU.

PyTorch

As of this writing, you must install the Preview (Nightly) build to train the PyTorch model on Apple Silicon GPUs. This will be stable in the PyTorch v1.12 release.

conda install pytorch torchvision  -c pytorch-nightly

To train PyTorch models on GPUs on Apple Silicon, set Metal Performance Shaders (MPS) as the backend.

device = torch.device("mps")
print(f"Using {device} device")
# Using mps device

Final thoughts

The repo below has example TensorFlow and PyTorch notebooks showing how to train deep learning models in Apple Silicon. Leave a comment below if you have any challenges setting up PyTorch and TensorFlow on Apple Silicon.

TensorFlow resources

Object detection with TensorFlow 2 Object detection API

How to create custom training loops in Keras

How to build CNN in TensorFlow(examples, code, and notebooks)

How to build artificial neural networks with Keras and TensorFlow

The Complete Data Science and Machine Learning Bootcamp on Udemy is a great next step if you want to keep exploring the data science and machine learning field.

Follow us on LinkedIn, Twitter, GitHub, and subscribe to our blog, so you don't miss a new issue.