Tensorflow with Metal on a M1 Mac

Photo by Ali Mahmoudi on Unsplash

From Tensorflow 2.5, the tensorflow-metal PluggableDevice to accelerate training with Metal on Mac GPUs is avaliable. Here we’ll install Tensorflow using conda. M1 native conda Installation using miniforge can be found here.

Say you already have conda installed. It’s better to create a virtual environment when using python.

$ conda create -n tf_mac python=3.9.6

This command creates env name tf_mac with python version 3.9.6, I’ve used newest python version to test if I was having any compatible issues so far okay. You can change name and python version but keep in mind python version 3.8~3.9 is required to install tensorflow-mac.

Now activate newly created env and install tensorflow-mac dependencies using commdand:

$ conda activate tf_mac
$ conda install -c apple tensorflow-deps

Install Tensorflow

$ pip install tensorflow-macos

Install Tensorflow metal plugin

$ pip install tensorflow-metal

Now you’re all set to use tensorflow and metal accelation. You can test with your test code.

Init Plugin
Init Graph Optimizer
Init Kernel
Metal device set to: Apple M1
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
simple_rnn (SimpleRNN)       (None, 20)                520       
=================================================================
Total params: 520
Trainable params: 520
Non-trainable params: 0
_________________________________________________________________
None

As you can see from mine, Metal device set to: Apple M1