Customized Tensorflow for macOS

This post is mostly inspired by this story.

Motivation

Since Macbook Pro doesn’t come with a discrete GPU, and the official Tensorflow installer lacks the flexibility for customizing. With this tutorial, you can unleash the full capability of Tensorflow with AVX2 and FMA enabled on macOS.

System Information

macOS Catalina 10.15.6
python 3.7.7
Tensorflow (CPU only) 1.14.1
Bazel 0.24.1

Method 1 (Didn’t try.)

Download pre-build Tensorflow installation package supporting AVX, AVX2, FMA and etc.
https://github.com/lakshayg/tensorflow-build

Method 2 (It works.)

Build your own installer from scratch!
Reason 1: If there is no matching version for your setups.
Reason 2: Free CPU stress test🔥

Step 1 Install Tensorflow pip package

pip install -U --user pip six numpy wheel setuptools mock 'future>=0.17.1'
pip install -U --user keras_applications --no-deps
pip install -U --user keras_preprocessing --no-deps

Step 2 Clone Tensorflow from Github

mkdir ~/Desktop/temp
cd ~/Desktop/temp
git clone https://github.com/tensorflow/tensorflow
cd tensorflow
git checkout r1.14.1

Step 3 Install Xcode

Go to app Store and search for Xcode. Once done installing, enter

sudo xcodebuild -license accept

to accept the license.

Step 4 Install Bazel

export BAZEL_VERSION=0.24.1
curl -fLO "https://github.com/bazelbuild/bazel/releases/download/${BAZEL_VERSION}/bazel-${BAZEL_VERSION}-installer-darwin-x86_64.sh"

Change the permission

chmod +x bazel-0.24.1-installer-darwin-x86_64.sh
./bazel-0.24.1-installer-darwin-x86_64.sh --user

Setup the path

export PATH="$PATH:$HOME/bin"

You can check if the installation is successful by entering

bazel version

and it should show the version.

Step 5 Setup customization options for Tensorflow

cd ~/Desktop/temp/tensorflow
./configure

You should specify the location of your python interpreter and library. If you are building this in an Anaconda environment, you can go with the default path by pressing Enter.

(tfavx) me@host /Users/thomas/Desktop/temp/tensorflow$ ./configure
WARNING: Running Bazel server needs to be killed, because the startup options are different.
WARNING: --batch mode is deprecated. Please instead explicitly shut down your Bazel server using the command "bazel shutdown".
You have bazel 0.24.1 installed.
Please specify the location of python. [Default is /Users/thomas/opt/anaconda3/envs/tfavx/bin/python]
Found possible Python library paths:
/Users/thomas/opt/anaconda3/envs/tfavx/lib/python3.7/site-packages
Please input the desired Python library path to use.  Default is [/Users/thomas/opt/anaconda3/envs/tfavx/lib/python3.7/site-packages]
Do you wish to build TensorFlow with XLA JIT support? [y/N]: Y
XLA JIT support will be enabled for TensorFlow.
Do you wish to build TensorFlow with OpenCL SYCL support? [y/N]: N
No OpenCL SYCL support will be enabled for TensorFlow.
Do you wish to build TensorFlow with ROCm support? [y/N]: N
No ROCm support will be enabled for TensorFlow.
Do you wish to build TensorFlow with CUDA support? [y/N]: N
No CUDA support will be enabled for TensorFlow.
Do you wish to download a fresh release of clang? (Experimental) [y/N]: N
Clang will not be downloaded.
Do you wish to build TensorFlow with MPI support? [y/N]: N
No MPI support will be enabled for TensorFlow.
Please specify optimization flags to use during compilation when bazel option "--config=opt" is specified [Default is -march=native -Wno-sign-compare]:
Would you like to interactively configure ./WORKSPACE for Android builds? [y/N]: N
Not configuring the WORKSPACE for Android builds.
Do you wish to build TensorFlow with iOS support? [y/N]: N
No iOS support will be enabled for TensorFlow.
Preconfigured Bazel build configs. You can use any of the below by adding "--config=<>" to your build command. See .bazelrc for more details.
--config=mkl          # Build with MKL support.
--config=monolithic   # Config for mostly static monolithic build.
--config=gdr          # Build with GDR support.
--config=verbs        # Build with libverbs support.
--config=ngraph       # Build with Intel nGraph support.
--config=numa         # Build with NUMA support.
--config=dynamic_kernels # (Experimental) Build kernels into separate shared objects.
Preconfigured Bazel build configs to DISABLE default on features:
--config=noaws        # Disable AWS S3 filesystem support.
--config=nogcp        # Disable GCP support.
--config=nohdfs       # Disable HDFS support.
--config=noignite     # Disable Apache Ignite support.
--config=nokafka      # Disable Apache Kafka support.
--config=nonccl       # Disable NVIDIA NCCL support.
Configuration finished

Step 6 Using Bazel to compile your Tensorflow installer

Before compiling, you might want to modify some codes. During my first attamp to compile, the process fails after compiling for 50 minutes. The part of error is as follows

tensorflow/python/lib/core/bfloat16.cc:682:8: error: no matching function for call to object of type '(lambda at tensorflow/python/lib/core/bfloat16.cc:637:25)'
if (!register_ufunc("greater_equal",CompareUFunc<Bfloat16GeFunctor>,
     ^~~~~~~~~~~~~~

To be honest, I don’t know what the hell is going on. After some googling, I find someone has a similar problem, and he provides a patch. Modify the file.

~/Desktop/temp/tensorflow/tensorflow/python/lib/core/bfloat16.cc

The line number may not be exactly the same. You should search for the function BinaryUFunc add const to the 2nd and 3rd arguments. Also, add a new template CompareUFunc with the corresponding argument type as BinaryUFunc. You can follow the screenshot above.

There is probably a better way to do this than manually modify the file. But it just works, we can begin our compilation.

These options enable AVX, AVX2 and FMA

bazel build -c opt --copt=-mavx --copt=-mavx2 --copt=-mfma --copt=-msse4.2 //tensorflow/tools/pip_package:build_pip_package

Well done! Now you can sit back and enjoy your Macbook taking off. It takes around 1 hour and 20 minutes for my Macbook Pro 2020 (2 GHz Intel Core i5) to finish the compilation.

Step 7 Building Wheel

./bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg

Now under /tmp/tensorflow_pkg/ you should find the installer

tensorflow-1.14.1-cp37-cp37m-macosx_10_15_x86_64.whl

Step 8 Install with PIP

pip install /tmp/tensorflow_pkg/tensorflow-1.14.1-cp37-cp37m-macosx_10_15_x86_64.whl

Step 9 Finish 🥰

Despite the power of AVX2 and FMA, it is still way slower than Tensorflow with GPU. But hey, something is better than nothing.

References
1. 中文版（版本不太一樣）
2. Tensorflow從原始碼開始建構
3. Installing Bazel on macOS