在macOS上加速pytroch的训练

Metal acceleration

PyTorch 使用新的 Metal Performance Shaders (MPS) 后端为 GPU 训练加速。MPS 后端扩展了 PyTorch 框架，提供了在 Mac 上设置和运行操作的脚本和功能。MPS 框架通过针对每个 Metal GPU 系列的独特特性进行微调的内核来优化计算性能。新的 mps 设备将机器学习计算图形和基元映射到 MPS Graph 框架和 MPS 提供的经过调整的内核上。

Requirements

• Mac computers with Apple silicon or AMD GPUs
• macOS 12.3 or later
• Python 3.7 or later
• Xcode command-line tools: xcode-select --install

Get started

可以使用 Anaconda 或 pip。请注意，苹果芯片 Mac 和英特尔 x86 Mac 的环境设置会有所不同。

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 Miniconda3-latest-MacOSX-x86_64.sh

pip

您可以使用 macOS 预装的 "pip3"。或者，你也可以从 Python 网站或 Homebrew 软件包管理器中安装。

2. Install

Anaconda

conda install pytorch torchvision torchaudio -c pytorch-nightly

pip

pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu

Building from source

构建支持 MPS 的 PyTorch 需要 Xcode 13.3.1 或更高版本。您可以在 Mac App Store 上下载最新的公开 Xcode 版本，或在 Mac App Store 上下载最新的测试版，或在 Apple Developer website 上下载最新的测试版。USE_MPS` 环境变量控制 PyTorch 的构建，并包含 MPS 支持。

要构建 PyTorch，请遵循 PyTorch 网站 上提供的说明。

3. Verify

You can verify mps support using a simple Python script:

import torch
if torch.backends.mps.is_available():
    mps_device = torch.device("mps")
    x = torch.ones(1, device=mps_device)
    print (x)
else:
    print ("MPS device not found.")

The output should show:

tensor([1.], device='mps:0')