Installing TensorFlow 2.5 and Jupyter Lab on Mac with M1

Last month, I finally painstakingly installed TensorFlow 2.4 and Jupyter Lab on my Mac with M1 (see the blog post). It worked nicely: 10 times faster than Colab, but also had a few issues like working only with Python 3.8, having to manually downgrade some packages such as numpy and h5py. Two days ago, Apple released TensorFlow 2.5 for Mac.

On https://github.com/apple/tensorflow_macos, they now redirect to TensorFlow 2.5, customized to run on the GPU instead of the CPU:

You can now leverage Apple’s tensorflow-metal PluggableDevice in TensorFlow v2.5 for accelerated training on Mac GPUs directly with Metal. Learn more here.

The page shows how to install it, and this time it seems much simpler.

Tensorflow Plugin - Metal - Apple Developer
Find presentations, documentation, sample code, and resources for building macOS, iOS, and tvOS apps with the Metal framework.

The steps are:

  • create a conda environment: conda create --name tf25 followed by conda activate tf25
  • install tensorflow-deps from the apple channel: conda install -c apple tensorflow-deps
  • install tensorflow-macos: pip install tensorflow-macos
  • install tensorflow-metal: pip install tensorflow-metal
  • install Jupyter Lab: pip install jupyterlab

Here's the output from installing tensorflow-deps.

(base) % conda create --name tf25
(base) % conda activate tf25
(tf25) % conda install -c apple tensorflow-deps
Collecting package metadata (current_repodata.json): done
Solving environment: failed with initial frozen solve. Retrying with flexible solve.
Solving environment: failed with repodata from current_repodata.json, will retry with next repodata source.
Collecting package metadata (repodata.json): done
Solving environment: done

## Package Plan ##

  environment location: /Users/anhtuan/miniforge3/envs/tf25

  added / updated specs:
    - tensorflow-deps


The following packages will be downloaded:

    package                    |            build
    ---------------------------|-----------------
    absl-py-0.10.0             |     pyhd8ed1ab_1          95 KB  conda-forge
    ca-certificates-2021.5.30  |       h4653dfc_0         136 KB  conda-forge
    certifi-2021.5.30          |   py39h2804cbe_0         142 KB  conda-forge
    flatbuffers-1.12.1         |       hbdafb3b_0        1007 KB  conda-forge
    grpcio-1.34.1              |   py39h1eaaf2b_0         1.8 MB  conda-forge
    h5py-3.1.0                 |nompi_py39h99babb8_100         1.1 MB  conda-forge
    libcurl-7.77.0             |       h8fe1914_0         301 KB  conda-forge
    libprotobuf-3.17.2         |       hccf11d3_0         2.1 MB  conda-forge
    numpy-1.19.5               |   py39h69a04d8_1         4.7 MB  conda-forge
    pip-21.1.2                 |     pyhd8ed1ab_0         1.1 MB  conda-forge
    protobuf-3.17.2            |   py39hfb83b0d_0         326 KB  conda-forge
    python-3.9.4               |h5b20da3_0_cpython        12.0 MB  conda-forge
    scipy-1.6.3                |   py39h5060c3b_0        18.1 MB  conda-forge
    tensorflow-deps-2.5.0      |                1           3 KB  apple
    wheel-0.35.1               |     pyh9f0ad1d_0          29 KB  conda-forge
    wrapt-1.12.1               |   py39h46acfd9_3          45 KB  conda-forge
    ------------------------------------------------------------
                                           Total:        43.0 MB

The following NEW packages will be INSTALLED:

  absl-py            conda-forge/noarch::absl-py-0.10.0-pyhd8ed1ab_1
  astunparse         conda-forge/noarch::astunparse-1.6.3-pyhd8ed1ab_0
  c-ares             conda-forge/osx-arm64::c-ares-1.17.1-h27ca646_1
  ca-certificates    conda-forge/osx-arm64::ca-certificates-2021.5.30-h4653dfc_0
  cached-property    conda-forge/noarch::cached-property-1.5.2-hd8ed1ab_1
  cached_property    conda-forge/noarch::cached_property-1.5.2-pyha770c72_1
  certifi            conda-forge/osx-arm64::certifi-2021.5.30-py39h2804cbe_0
  flatbuffers        conda-forge/osx-arm64::flatbuffers-1.12.1-hbdafb3b_0
  gast               conda-forge/noarch::gast-0.4.0-pyh9f0ad1d_0
  google-pasta       conda-forge/noarch::google-pasta-0.2.0-pyh8c360ce_0
  grpcio             conda-forge/osx-arm64::grpcio-1.34.1-py39h1eaaf2b_0
  h5py               conda-forge/osx-arm64::h5py-3.1.0-nompi_py39h99babb8_100
  hdf5               conda-forge/osx-arm64::hdf5-1.10.6-nompi_h0fc092c_1114
  keras-preprocessi~ conda-forge/noarch::keras-preprocessing-1.1.2-pyhd8ed1ab_0
  krb5               conda-forge/osx-arm64::krb5-1.19.1-hd92b7a7_0
  libblas            conda-forge/osx-arm64::libblas-3.9.0-9_openblas
  libcblas           conda-forge/osx-arm64::libcblas-3.9.0-9_openblas
  libcurl            conda-forge/osx-arm64::libcurl-7.77.0-h8fe1914_0
  libcxx             conda-forge/osx-arm64::libcxx-11.1.0-h168391b_0
  libedit            conda-forge/osx-arm64::libedit-3.1.20191231-hc8eb9b7_2
  libev              conda-forge/osx-arm64::libev-4.33-h642e427_1
  libgfortran        conda-forge/osx-arm64::libgfortran-5.0.0.dev0-11_0_1_hf114ba7_22
  libgfortran5       conda-forge/osx-arm64::libgfortran5-11.0.1.dev0-hf114ba7_22
  liblapack          conda-forge/osx-arm64::liblapack-3.9.0-9_openblas
  libnghttp2         conda-forge/osx-arm64::libnghttp2-1.43.0-hf3018f0_0
  libopenblas        conda-forge/osx-arm64::libopenblas-0.3.15-openmp_hf330de4_1
  libprotobuf        conda-forge/osx-arm64::libprotobuf-3.17.2-hccf11d3_0
  libssh2            conda-forge/osx-arm64::libssh2-1.9.0-hb80f160_6
  llvm-openmp        conda-forge/osx-arm64::llvm-openmp-11.1.0-hf3c4609_1
  ncurses            conda-forge/osx-arm64::ncurses-6.2-h9aa5885_4
  numpy              conda-forge/osx-arm64::numpy-1.19.5-py39h69a04d8_1
  openssl            conda-forge/osx-arm64::openssl-1.1.1k-h27ca646_0
  opt_einsum         conda-forge/noarch::opt_einsum-3.3.0-pyhd8ed1ab_1
  pip                conda-forge/noarch::pip-21.1.2-pyhd8ed1ab_0
  protobuf           conda-forge/osx-arm64::protobuf-3.17.2-py39hfb83b0d_0
  python             conda-forge/osx-arm64::python-3.9.4-h5b20da3_0_cpython
  python_abi         conda-forge/osx-arm64::python_abi-3.9-1_cp39
  readline           conda-forge/osx-arm64::readline-8.1-hedafd6a_0
  scipy              conda-forge/osx-arm64::scipy-1.6.3-py39h5060c3b_0
  setuptools         conda-forge/osx-arm64::setuptools-49.6.0-py39h2804cbe_3
  six                conda-forge/noarch::six-1.15.0-pyh9f0ad1d_0
  sqlite             conda-forge/osx-arm64::sqlite-3.35.5-hc49ca36_0
  tensorflow-deps    apple/osx-arm64::tensorflow-deps-2.5.0-1
  termcolor          conda-forge/noarch::termcolor-1.1.0-py_2
  tk                 conda-forge/osx-arm64::tk-8.6.10-hf7e6567_1
  typing_extensions  conda-forge/noarch::typing_extensions-3.7.4.3-py_0
  tzdata             conda-forge/noarch::tzdata-2021a-he74cb21_0
  wheel              conda-forge/noarch::wheel-0.35.1-pyh9f0ad1d_0
  wrapt              conda-forge/osx-arm64::wrapt-1.12.1-py39h46acfd9_3
  xz                 conda-forge/osx-arm64::xz-5.2.5-h642e427_1
  zlib               conda-forge/osx-arm64::zlib-1.2.11-h31e879b_1009


Proceed ([y]/n)? y

It looks like they also downgraded to numpy 1.19.5 as I found out I had to do with TF2.4. It's nice to know that it's the official way to set it up though, instead of blindly following the tip from some obscure StackOverflow posts.

After installing everything, I ran Jupyter Lab (jupyter lab) and ran some training. While the TF2.4 used to show 140% CPU activity on the Python 3.8 process in the activity monitor, now TF 2.5 showed up as Python 3.9 process which does 90% GPU activity.

And performance wise, training takes half as long as it used to. WOW!

TensorFlow now uses my Mac's GPU