Install PyTorch 2.0 GPU/MPS for Mac M1/M2 with Conda

In this post, we will look at how to install PyTorch 2.0 from the beginning on a Mac M1/M2 Apple silicon and set it up in a Conda environment. We do everything through Conda and Jupyter Notebook.

PyTorch Training on Mac

PyTorch 2.0 support GPU-accelerated training on Mac. Until now, PyTorch training on Mac only leveraged the CPU, but with the upcoming PyTorch release, you can take advantage of Apple silicon GPUs for significantly faster model training.

Apple Metal Performance Shaders (MPS)

The MPS backend extends the PyTorch framework by providing scripts and capabilities to set up and run operations on Mac. MPS optimizes compute performance with kernels that are fine-tuned for the unique characteristics of each Metal GPU family.

Requirements

Mac computers with Apple silicon or AMD GPUs
macOS 12.3 or later
Python 3.7 or later

Install Miniconda Apple M1/M2

In this post we are going to use Miniconda, it’s a Python environment and it has a lot of scientific packages available for data science. I use Miniconda rather than Anaconda they’re both from the same company but Miniconda does not install a whole plethora of additional packages.

Anaconda directly supports Windows, Mac, and Linux. Make sure that you select the Miniconda version that corresponds to your operating system. It is particularly important to choose M1/M2 Metal if you have a later (non-Intel) Mac.

If you have a Mac and wish to use M1/M2 MPS make sure to install the ARM64 version of Miniconda. Miniconda is the minimal set of features from the extensive Anaconda Python distribution. Download Miniconda from this URL.

Apple silicon

curl -O https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh
sh Miniconda3-latest-MacOSX-arm64.sh

x86

curl -O https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh
sh Miniconda3-latest-MacOSX-x86_64.sh

If you have an Intel x86 version of Miniconda or Anaconda installed on your computer and you try to use the Apple silicon on top of it that’s not going to work you need an arm64 version of Miniconda or Anaconda.

If you’re using zsh and it has not been set up to read .bashrc, you need to add the Miniconda directory to the zsh shell path environment variable. Add this to your .zshrc:

export PATH="/home/username/miniconda/bin:$PATH"

Make sure to replace /home/username/miniconda with your actual path. Save, exit the terminal, and then reopen the terminal. conda command should work.

Here’s the check to type Python.

Conda check Platform — if you’ve got x86 it gonna not work.

Install Jupyter Notebook

Now that Miniconda successfully installed, the next thing we’re going to do is install Jupyter Notebook. We want to install Jupyter Notebook to the base environment and then it will use any other environments that we create underneath.

conda install -y jupyter

This instruction is talking about how to install this with an M1/M2 Mac but you could install this CPU only. Nvidia cuda GPU is not going to work on a Mac.

You must make sure that PyTorch has the version of Python that it is compatible with. The best way to accomplish this is with an Anaconda environment. Each environment that you create can have its own Python version, drivers, and Python libraries. I suggest that you create an environment to hold the Python instance. Use the following command to create your environment. I am calling the environment torchEnv.

conda create --name torchEnv

It’s going to install Python for us and a few bunch of other things there are pip setup tools so to activate that environment we could use this full command.

Register Conda Environment

The following command registers your PyTorch environment. Again, make sure you activate your new PyTorch environment.

conda activate torchEnv

conda install -c anaconda ipykernel
python -m ipykernel install --user --name=torchEnv

We’re just active conda environment that you just created and registered as a Jupyter Kernal If you don’t do this it won’t show up on your list of available Kernels in Jupyter Notebook.

Install PyTorch

What we’re going to do in this post is set up a Conda base environment for data science and machine learning on Apple silicon with PyTorch. Install the PyTorch 2.0+ version for Mac.

# MPS acceleration is available on MacOS 12.3+
conda install pytorch::pytorch torchvision torchaudio -c pytorch

We’re going to Install common data science packages such as Jupyter, Numpy, Pandas, Matplotlib, tqdm, and Sciket-learn just a handful of helpful data science libraries.

conda install jupyter pandas numpy matplotlib scikit-learn tqdm

Run Jupyter Notebook

Start a Jupiter server using run command, this should launch a Jupiter notebook server in our browser.

[(base) [email protected] ~ % jupyter notebook

First thing you want to do is make sure that you’ve got the right Kernel running . You can see that I installed that torchEnv and it shows up here.

We’re just going to run this code just to verify that things are indeed working.The target devices is MPS.

import torch
import platform

has_gpu = torch.cuda.is_available()
has_mps = getattr(torch,'has_mps',False)

device = "mps" if getattr(torch,'has_mps',False) \
    else "gpu" if torch.cuda.is_available() else "cpu"

print(f"Python Platform: {platform.platform()}") #Python Platform: macOS-13.3.1-arm64-arm-64bit
print(f"PyTorch Version: {torch.__version__}") #PyTsorch Version: 2.0.1
 
print("GPU is", "available" if has_gpu else "NOT AVAILABLE") GPU is NOT AVAILABLE
print("MPS is", "AVAILABLE" if has_mps else "NOT AVAILABLE") #MPS is AVAILABLE

print(f"Target device is {device}") #Target device is mps