Device-agnostic code and using Mac M1 and M2 Chips on TensorFlow

Ata_Tekeli · June 26, 2023, 10:09pm

Hi,

Even though I tried to look at documentation and tutorials on the net, I failed to find a good tutorial or documentation on this topic. How can I create a device agnostic code so that I can use Mac M1 GPU easily?

Yingding · June 29, 2023, 9:04pm

The code for TensorFlow is mostly device agnostic. The underline framework of Tensorflow is taking care of the device part for you. Please take this example on M1.

github.com

yingding/tf-every-day/blob/main/10zenml/requirements.txt

zenml==0.39.1 # 0.36.1 no fastapi dependency
tensorflow-macos==2.12.0
tensorflow-metal==0.8.0
pandas==2.0.1 #1.5.3
scikit-learn==1.2.2
fastapi==0.95.2 # need by zenml onwards 0.36.0
# model-card-toolkit
# model-card-toolkit==2.0.0 # still missing tensorflow-data-validation for generating data statistics

You just need to use tensorflow-macos and tensorflow-metal packages.

The same tensorflow code runs on different devices.

github.com

yingding/tf-every-day/blob/main/10zenml/dist_train.py

## Test for distributed training
## https://www.tensorflow.org/guide/distributed_training

import tensorflow as tf

print(tf.version.VERSION)

strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0"], cross_device_ops=tf.distribute.NcclAllReduce())
# strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")

with strategy.scope():
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(1, input_shape=(1,))
    ])
    # it shoud be inside the strategy scope
    model.compile(loss='mae', optimizer='sgd')

dataset = tf.data.Dataset.from_tensors(([1.], [1.]))

# https://stackoverflow.com/questions/72740907/tensorflow-cant-apply-sharing-policy-file-when-using-mirrored-strategy

This file has been truncated. show original

As I stated mostly, because the device or accelerator ID need to be changed in case of distributed training, but most case on a single host no further code modification are necessary.

Hope this can be helpful for you.