Get started with tensorflow-metal
Accelerate the training of machine learning models with TensorFlow right on your Mac. Install base TensorFlow and the tensorflow-metal PluggableDevice to accelerate training with Metal on Mac GPUs.
Learn about TensorFlow PluggableDevices
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
arm64 : Apple silicon
bash ~/miniconda.sh -b -p $HOME/miniconda
source ~/miniconda/bin/activate
conda install -c apple tensorflow-depsx86 : AMD
Virtual environment
python3 -m venv ~/venv-metal
source ~/venv-metal/bin/activate
python -m pip install -U pip2. Install base TensorFlow
python -m pip install tensorflow-macos3. Install tensorflow-metal plug-in
python -m pip install tensorflow-metal4. Verify
You can verify using a simple script:
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=True)
model.compile(optimizer="adam", loss=loss_fn, metrics=["accuracy"])
model.fit(x_train, y_train, epochs=5, batch_size=64)