Pytorch for Mac M1/M2 with GPU acceleration 2023. Jupyter and VS Code setup for PyTorch included.
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.
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.
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
- Python Environments: Display all the python environment including virtual environments.
- 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.