Get Started with TensorFlow-Metal on Mac

Accelerate the training of machine learning models with TensorFlow right on your Mac. This guide will walk you through installing the base TensorFlow and the tensorflow-metal PluggableDevice to harness the power of Metal on Mac GPUs.

Requirements

  • Mac computers with Apple silicon or AMD GPUs
  • macOS 12.0 or later (Get the latest beta)
  • Python 3.8 or later
  • Xcode command-line tools: xcode-select --install

Get Started

1. Set up the Environment

First, create a virtual environment to keep your dependencies isolated.

python3.10 -m venv ~/venv-metal
source ~/venv-metal/bin/activate
python -m pip install -U pip

2. Install Base TensorFlow

Depending on your TensorFlow version, you will install it as follows:

  • For TensorFlow version 2.13 or later:

    python3 -m pip install tensorflow
    
  • For TensorFlow version 2.12 or earlier:

    python3 -m pip install tensorflow-macos
    
  • Requirement (install bolth of them)

    3. Install tensorflow-metal Plug-in

Install the tensorflow-metal plug-in to enable GPU acceleration.

python3 -m pip install tensorflow-metal

4. Verify the Installation

You can verify your installation with a simple script. This script trains a ResNet50 model on the CIFAR-100 dataset.

import tensorflow as tf

cifar = tf.keras.datasets.cifar100
(x_train, y_train), (x_test, y_test) = cifar.load_data()
model = tf.keras.applications.ResNet50(
    include_top=True,
    weights=None,
    input_shape=(32, 32, 3),
    classes=100,)

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False)
model.compile(optimizer="adam", loss=loss_fn, metrics=["accuracy"])
model.fit(x_train, y_train, epochs=5, batch_size=64)

Troubleshooting

Common Errors

  • Error: “Could not find a version that satisfies the requirement tensorflow-macos (from versions: none).”
    • Ensure the Python version used in the environment is supported (Python 3.8, 3.9, 3.10).
  • Error: “No registered OpKernel. (OpKernel was found, but attributes didn’t match) Requested Attributes: dtype=DT_COMPLEX64.”
    • Complex data type isn’t supported by tensorflow-metal.
  • Error: “Cannot assign a device for operation: Could not satisfy explicit device specification because the node was colocated with a group of nodes that required incompatible device.”
    • A colocation issue takes place when an operation doesn’t have a GPU implementation available. Report missing operations on the Apple Developer Forums.

Performance Considerations

  • CPU faster than GPU:
    • This can occur on small networks with small batch sizes due to overhead in dispatching computations to the GPU. For larger batches or models, the GPU will better leverage parallelism for improved performance.

Resources

For details on versioning and additional notes on notable implemented features, refer to the table in the PyPI repository.

To ask questions and share feedback, visit the Apple Developer Forums.

Currently Not Supported

  • Multi-GPU support
  • Acceleration for Intel GPUs
  • TensorFlow V1 networks