Pytorch for Mac M1/M2 with GPU acceleration 2023. Jupyter and VS Code setup for PyTorch included.

Image Source: https://wandb.ai/

Introduction

Among the numerous deep learning frameworks available, PyTorch stands tall as a powerful and versatile platform for building cutting-edge machine learning models. Its ability to leverage GPU acceleration has undoubtedly been a game-changer, unleashing the potential of modern hardware and drastically speeding up computation.

For Mac users wielding the mighty M1 or M2 chips, tapping into the full potential of PyTorch with GPU acceleration can be a transformative experience. These ARM-based Apple chips have set new benchmarks in performance and energy efficiency, making them a favourite among developers and AI enthusiasts alike.

In this comprehensive guide, we embark on an exciting journey to unravel the mysteries of installing PyTorch with GPU acceleration on Mac M1/M2 along with using it in Jupyter notebooks and VS Code.

Before we begin, make sure you have your seatbelts fastened, your Mac powered up, and your enthusiasm fueled. Let’s embark on this exhilarating adventure, where we’ll pave the way for unleashing the full potential of PyTorch with GPU acceleration on your Mac M1/M2 machine!

Pytorch Metal Performance Shader(MPS)

Every Apple silicon Mac has a unified memory architecture, providing the GPU with direct access to the full memory store. This makes Mac a great platform for machine learning, enabling users to train larger networks or batch sizes locally. This reduces costs associated with cloud-based development or the need for additional local GPUs. The Unified Memory architecture also reduces data retrieval latency, improving end-to-end performance.

MPS basically allows you to get access to Apple GPU while training which allows faster training.

Accelerated GPU training and evaluation speedups over CPU-only (times faster)

The above chart shows the comparison of training and evaluation between the CPU and GPU . Source: https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/

Note: It only works with Apple silicon macs and not with intel powered macs.

Implementation

Steps for installing PyTorch within conda environment

Step 1: Create a new conda environment

conda create --name ENV_NAME python=3.9

Step 2: Activate the conda enviroment

conda activate ENV_NAME

Step 3: Install PyTorch

Paste the below code in the terminal (make sure you’re in the env)

conda install pytorch::pytorch torchvision torchaudio -c pytorch

If you’re using pip instead of conda or you would like to use C++/Java as a language i encourage you to go to pytorch website and according to your preferences copy the generated command and paste into the terminal.

PyTorch

Congratulations PyTorch is now installed!

Running PyTorch in Jupyter Notebook

Step 1: Install Jupyter Notebook

conda install -c conda-forge notebook
conda install -c conda-forge nb_conda_kernels

Step 2: Install Jupyter Lab

conda install -c conda-forge jupyterlab
conda install -c conda-forge nb_conda_kernels

Step 3: Start Jupyter Notebook

jupyter notebook

Jupyter notebook will be opened. Now create a new file or open an exisiting one and type the following command to check whether the pytorch is correctly installed and the version of it.

import torch

print(torch.__version__)

It should output something like this

So the 2.0.1 here states the version.

Check for the MPS(Metal Performance Shader) i.e Apple Metal GPU

# Is MPS even available? macOS 12.3+
print(torch.backends.mps.is_available())

# Was the current version of PyTorch built with MPS activated?
print(torch.backends.mps.is_built())

If the above code outputs “true” & “true” then it’s time to celebrate because you now have access to Apple Metal GPU.

Running PyTorch in VS Code

Step 1: Activate your conda environment and then type following command in the terminal(make sure you have vs code installed) this will open VS code .

code . 

Now make a new file with .ipynb extension.

As soon as you create a new file you may need to select a kernel. I will be attaching photos which will guide you how to do so.

Step 2: Select kernel

After selecting kernel chances are your conda environment will be displayed click the environment and you’re ready to go but if not follow the next step.

Step 3: Click on select another kernel which will display two options

1. Python Environments: Display all the python environment including virtual environments.
2. Existing jupyter server: By providing the URL of running jupyter notebook you can access the notebook in the vs code.

Step 4: Select Python Environments. It would display list of environments

Choose your environment and tadaaa!!! All the libraries and frameworks will be loaded present in that specific environment in our case pytorch.

Now check for the PyTorch version.

import torch

print(torch.__version)

It may output pytorch version.

Check for the MPS(Metal Performance Shader) i.e Apple Metal GPU

# Is MPS even available? macOS 12.3+
print(torch.backends.mps.is_available())

# Was the current version of PyTorch built with MPS activated?
print(torch.backends.mps.is_built())

If the above code outputs “true” & “true” then it’s time to celebrate because you now have access to Apple Metal GPU.

I hope you find this article helpful. Support my work by following me.