How To Use Apple’s TensorFlow Fork and Jupyter with Visual Studio Code on Apple Silicon

I bought the latest M1 MacBook Pro as soon as Apple announced them in 2020. I’ve been a great fan of Apple’s work on the arm64 architecture, and I couldn’t wait to try out the chip that was screaming in my iPad Pro on a laptop. One of the interesting consequences of Apple moving to the arm64 architecture for macOS is its impact on how you can access the hardware resources for machine learning (ML) and artificial intelligence. On Intel Macs you could easily run Python 3, scikit-learn, matplotlib, tensorflow and jupyterlab, etc. without much hassle and access the onboard AMD or Nvidia GPU. This was great for learning ML and training models to see if your feature design was successful before moving to a cloud environment like Azure ML.

Apple Silicon for Machine Learning

Now that Apple switched to their own Apple Silicon the GPU and CPU is now integrated. The performance improvements Apple has managed in their integrated SoC is revolutionary. Developers will want to use the full power of these chips. This effectively creates a new stack that the ML ecosystem will need to support. To help with this, Apple released a new developer API called ML Compute during WWDC 2020. Complimenting their previously released Core ML framework, the primary benefit of using ML Compute is to expose the BNNS primitives made available by the Accelerate framework for the CPU and Metal Performance Shaders for the GPU in Apple Silicon directly in a low level API. In other words, with ML Compute you can train TensorFlow models right on the Mac and take advantage of accelerated CPU and GPU training on Apple Silicon without first having to export the model and supply it to the Core ML framework. This will also help Google’s attempt at moving TensorFlow to mobile devices with TensorFlow Lite. A lot of impressive charts can be seen in the blog post Apple made announcing ML Compute.

How To Set Up Your Apple Silicon Mac To Use TensorFlow and Jupyter

It can be assumed TensorFlow was what Apple was most targeting with ML Compute, and they have done most of the homework for the ML community by releasing a fork of TensorFlow that uses ML Compute. Running this script will create a virtual environment in Python using the Python3 included with XCode Command Line Tools. The script will do all of this for you.

Once you have your virtual environment created, it is a good time to launch it and begin loading the rest of the toolkit you want. I’ve had difficulty running pip from the Visual Studio Code Python Terminal, so I recommend doing them all from the Terminal.

Here are some examples of packages to install so you can run the Jupyter notebooks:

pip install --upgrade -t "${VIRTUAL_ENV}/lib/python3.8/site-packages/" ipykernel
pip install --upgrade -t "${VIRTUAL_ENV}/lib/python3.8/site-packages/" matplotlib

How To Set Up Your Visual Studio Code to Use The TensorFlow Virtual Environment

After you have your Python virtual environment set up, you will need to tell Visual Studio Code how to use the TensorFlow environment you just created. What is nice about the Python installed by Apple in XCode Command Line Tools is that it works well with Visual Studio Code, unlike the version bundled with macOS.

Since the TensorFlow environment has a lot of path based configuration required, you cannot just select the Python binary from Visual Studio Code environment. Instead, you will need to tell VS Code where to look for virtual environments and have it do the work of loading it in the list.

  1. Launch Visual Studio Code.
  2. Type ⌘+Shift+P to bring up the Command Palette.
  3. Select Your User Settings by Typing “Preferences: Open User Settings

4. Go to the python.venvPath by putting that string in the search box and put the location where you installed the TensorFlow Virtual Environment.

5. Reboot Visual Studio Code. In the list of Python environments available, you’ll now see the TensorFlow VM that Apple created!

 

Once selected, your code should execute in the TensorFlow virtual environment. You can check this by looking at the Python environment listed at the bottom of Visual Studio Code.

Congratulations! You now have a Tensor Flow virtual environment running on Apple Silicon and able to run your models using Visual Studio Code and Jupyter.