How to install tensorflow with GPU for Apple Silicon and Windows with nVidia GPU

I have been spending time installing, got the GPU working, then re-installing and finding errors installing over and over again. I never learn. So I am hoping to document this done to help myself, or someone stumbled upon this article happen to have the similar issues.

Python Version Matters

So I am the kind of person that always install latest python and breaking so manythings. It turns out no matter what version python you install, tensorflow also installs, except it wont work if the versions are incompatible. Most likely you will not be using GPU while training, or bump into some weird errors/crashs. Tensorflow is always a couple of version behind. At the time of writing (July 2024), I am only able to get python working (for both Windows and Mac) with python 3.9. Although on tensorflow website, it says it supports 3.8–3.11, I was never able to get GPU working.

Use Conda

I like to use pip and unsure why we needed conda anyway. I was wrong. The reason is while both allowing user to create virtual environments, conda allows user to specify python versions, and based on previous experience, it is a lot easier to setup with latest version of python, then create different python environments to better suit your needs. Plus managing the environments for different project to share is also made easier.

Install Steps

Pretty much from the tensorflow installation guide.

1. Windows and Mac: Create a conda environment:

conda create --name tf_py3.9 python=3.9
conda activate tf_py3.9

2. Install Dependencies

For Windows Native:

• Install latest nVidia Graphics Driver
• Install cuda and cuDNN

conda install -c conda-forge cudatoolkit=11.2 cudnn=8.1.0

For MacOS (Apple Silicon):

• Install Xcode Command Line tools

xcode-select --install 

3. Install Tensorflow

For Windows Native:

pip install --upgrade pip
pip install "tensorflow<2.11" 

For Windows WSL2

pip install tensorflow[and-cuda]

For MacOS (Applie Silicon)

python -m pip install tensorflow

4. To test if the install is correct. (using GPU not CPU) Create a python shell and run below code.

# Ensure TensorFlow uses the GPU
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        # Currently, memory growth needs to be the same across GPUs
        for gpu in gpus:
            print(gpu)
            tf.config.experimental.set_memory_growth(gpu, True)
        print("TensorFlow is using the GPU.")
    except RuntimeError as e:
        print(e)
else:
    print("TensorFlow is not using the GPU. Check your TensorFlow installation.")

Conclusion

Although looked easy, it takes some time to setup correctly. The tensflow install guide is not the most user friendly. I’d say creating virtual environment with conda, then test out may be the best option to quickly find a working environment. I used the same method tested 3.8, 3.9, 3.10, 3.11, 3.12. The only one worked for me is 3.9 so far.

If followed the the steps and install is complete without errors, also GPU test is passed. Then grats, you can train your super model now. If not, don’t go down a rabbit hold to google on errors I’d say, check compatibility will help you better.