Training PyTorch models on a Mac M1 and M2

PyTorch models on Apple Silicon M1 and M2

Next article: Segment Anything Model (SAM) on Apple Silicon M1 and M2

Segment Anything Model (SAM) on Apple Silicon M1 and M2

Metal acceleration

PyTorch utilizes the Metal Performance Shaders (MPS) backend for accelerating GPU training, which enhances the framework by enabling the creation and execution of operations on Mac. The MPS backend is optimized for compute performance by utilizing kernels that are tailored to the specific characteristics of each Metal GPU family. Through the use of the MPS Graph framework and MPS-tuned kernels, the mps device can efficiently map machine learning computational graphs and primitives.

Requirements

• Apple Silicon Mac (M1, M2, M1 Pro, M1 Max, M1 Ultra, etc).
• macOS 12.3+ (PyTorch will work on previous versions but the GPU on your Mac won’t get used, this means slower code).
• macOS 12.3 or later
• Python 3.7 or later
• Xcode command-line tools: xcode-select --install

Get started

You have the option to utilize either Anaconda or pip for installation, although the process for setting up the environment will vary depending on whether you have a Mac with Apple silicon or Intel x86. To enable MPS device acceleration, access the PyTorch installation selector and select Preview (Nightly). MPS backend support is included in the official release of PyTorch 1.12. By opting for the Preview (Nightly) build of PyTorch, you will be able to take advantage of the most up-to-date mps support on your device.

1. Set up

Anaconda

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…