PyTorch for Mac M1/M2 with GPU Acceleration: A Small Guide

Mac users with M1 or M2 chips can now harness the power of PyTorch with GPU acceleration, thanks to the Metal Performance Shaders (MPS) framework. Apple’s ARM-based silicon provides a unified memory architecture, enabling the GPU to directly access memory, improving training speed and reducing costs compared to cloud-based alternatives.

Why PyTorch on Mac?

The M1 and M2 chips make Mac a compelling platform for machine learning due to their performance and energy efficiency. Utilizing PyTorch on these machines can significantly enhance training times by using Apple’s GPU instead of relying solely on the CPU. This approach is only compatible with Apple Silicon Macs, not Intel-powered ones.

Step-by-Step Guide to Installing PyTorch with GPU Acceleration

1. Set Up Your Environment:

Create a Conda Environment:

conda create --name ENV_NAME python=3.9

conda activate ENV_NAME

Install PyTorch:

conda install pytorch torchvision torchaudio -c pytorch

Alternatively, use the PyTorch website to find the correct command for pip or other languages.

2. Running PyTorch in Jupyter Notebook:

Install Jupyter Notebook:

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

Start Jupyter Notebook:

jupyter notebook

Test the installation by checking PyTorch’s version:

import torch
print(torch.version)

3. Running PyTorch in VS Code:

Open VS Code in Conda Environment:

code .

Create a new .ipynb file and select the appropriate Python kernel for your environment.

4. Verify MPS (Metal Performance Shaders) Availability:

print(torch.backends.mps.is_available())
print(torch.backends.mps.is_built())


## If both return "True," you are successfully using GPU acceleration with PyTorch on your Mac!

By following these steps, you’ll unlock the full potential of your Mac’s hardware for deep learning tasks using PyTorch. Enjoy the speed and efficiency that Apple’s GPU acceleration offers!