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