TensorFlow on Mac M1 GPU: Installation and Performance

I have written an article about installing and running PyTorch on Mac M1 GPU. This article is on TensorFlow.

To get started, the following Apple’s document would be useful: https://developer.apple.com/metal/tensorflow-plugin/

As I already have my environment set up and TensorFlow installed, I went straight away to install tensorflow-metal plug-in. I also installed RUST for Transformers.

python -m pip install tensorflow-metal
curl https://sh.rustup.rs -sSf | bash -s - -y - no-modify-path

I modified the script for verification to compare the performance of running TensorFlow on M1 and CPU. The task is categorical classification of CIFAR-100 images.

from time import process_time

def testgpu():
    t0 = process_time()
    import tensorflow as tf
    tf.config.list_physical_devices()
    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)
    t1 = process_time()
    print("Total time with gpu: ", t1-t0)

def testcpu():
    t0 = process_time()
    import tensorflow as tf
    with tf.device('/CPU:0')…