How to set up a good PyTorch environment on Apple Silicon
Apple silicon refers to the arm64 architecture chips designed by Apple and used by new Macbook laptops, also known as M1, M2 Macbooks.
It's different from old generations of Macbooks, because those use Intel x86_64 architecture chips. You can find out these information on your Macbook's system information.
So why do you care?
You care because you care about performance when you want to do AI, research, development, etc. It's best to get the best performance natively out of your Macbook instead of the other way round. So make sure you check whether you are installing a package built for arm64 or x86_64, before doing the real AI stuff.
Install arm64 miniconda
The first step to make things right is to install the Miniconda version for MacOSX arm64:
mkdir -p ~/miniconda3
curl https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh -o ~/miniconda3/miniconda.sh
bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
rm -rf ~/miniconda3/miniconda.sh
Install arm64 pytorch
Since we installed miniconda's arm64 version, it automatically installs the correct pytorch and python versions in the arm64 architecture for us, using the following command:
It installs the correct version for arm64, which you can see from the command's output, something like the this:
Verify MPS support in pytorch
MPS stands for Metal Performance Shader, which is Apple's GPU programming language, corresponding to Nvidia's CUDA. We need MPS available so as to do efficient ML computation on Macbook's Apple Silicon.
The code below should output no assertion errors if you are running it on Apple Silicon, as well as having installed the arm64 version of pytorch.
As of now, you have set up a good PyTorch environment that allows you to take adavantage of all the computational power Apple Silicon has to offer. Enjoy!