Using TensorFlow 2.8 on an Apple Silicon arm64 chip
Tags: tensorflow, apple, gpu, hardware
Published: 2024-06-06 06:23:14

My computer recently had an unfortunate interface with dihydrogen monoxide. To be determined if it will come back to life but it's not looking good. So, I bought a new Macbook which means, the M3 chip (arm64). I had a nice experience using the M1 from my previous job so I was looking forward to it.

Of course 😩   the `x86` vs `arm` architecture issues started immediately when I tried using TensorFlow.

Here's how I fixed it. The pull request: [deepcell-imaging#229](https://github.com/dchaley/deepcell-imaging/pull/229)

DeepCell uses TF 2.8 so that's what we have to use. Unfortunately the [2.8.4 package](https://pypi.org/project/tensorflow/2.8.4/#files) doesn't come with ARM binaries. Incidentally [TF 2.16.1](https://pypi.org/project/tensorflow/2.16.1/#files) *does* have arm64 binaries ... but I can't use it here 😑

Apple has [some documentation](https://developer.apple.com/metal/tensorflow-plugin/) for installing TensorFlow and the "metal" plugin. In particular,

> For TensorFlow version 2.12 or earlier:
> ```
> python -m pip install tensorflow-macos
> ```

In our case we need tensorflow-macos==2.8.0 as found in the [tensorflow-macos release history](https://pypi.org/project/tensorflow-macos/#history). Unfortunately the [files list](https://pypi.org/project/tensorflow-macos/2.8.0/#files) reveals there's no Python 3.10 distribution so I need to downgrade to Python 3.9.

As for tensorflow-metal the [package documentation](https://pypi.org/project/tensorflow-metal/) says we need v0.4.0.

I packaged a new requirements file for Mac arm64 users:

```
$ cat requirements-mac-arm64.txt
tensorflow-macos==2.8.0
tensorflow-metal==0.4.0
```

Then you install the mac requirements:

```bash
pip install -r requirements-mac-arm64.txt
```

Of course, the shenanigans don't stop! Running `pip install -r requirements.txt` fails to install DeepCell, because it depends on `tensorflow` – not `tensorflow-macos` (which provides the same Python module `tensorflow`).

So I ran it this way to skip dependencies after installing the ones we could:

```bash
pip install -r requirements-mac-arm64.txt
pip install -r requirements.txt
pip install -r requirements.txt --no-deps
```

Then I got an interesting protobuf failure.

```
If this call came from a _pb2.py file, your generated code is out of date and must be regenerated with protoc >= 3.19.0.
If you cannot immediately regenerate your protos, some other possible workarounds are:
 1. Downgrade the protobuf package to 3.20.x or lower.
 2. Set PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python (but this will use pure-Python parsing and will be much slower).

More information: https://developers.google.com/protocol-buffers/docs/news/2022-05-06#python-updates
```

Quick fix: grab the most recent 3.20.x protobuf version, [3.20.3](https://pypi.org/project/protobuf/3.20.3/).

Apple provides a test 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=False)
model.compile(optimizer="adam", loss=loss_fn, metrics=["accuracy"])
model.fit(x_train, y_train, epochs=5, batch_size=64)
```

One 180 MB model download later … we're golden.

```
2024-06-05 23:10:30.794862: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:113] Plugin optimizer for device_type GPU is enabled.
```

Just to confirm, let's check Activity Monitor – and yes, it's using the gpu. 🎉 😤

![Screenshot of the Activity Monitor showing GPU % usage and GPU time used.](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/pf6kqzmhplcq330hy92c.png)

Phew. Well, hopefully this is a one-time thing. Most of our development is cloud which is x86, the more common binary format.

Until our next adventure with binaries ✌


Cover image by <a href="https://unsplash.com/@karishea?utm_content=creditCopyText&utm_medium=referral&utm_source=unsplash">Kari Shea</a> on <a href="https://unsplash.com/photos/macbook-pro-on-top-of-brown-table-1SAnrIxw5OY?utm_content=creditCopyText&utm_medium=referral&utm_source=unsplash">Unsplash</a>