在Mac中加速PyTorch训练教程

一、Metal 加速

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

二、环境要求

  • 配备 Apple silicon 或 AMD GPU 的 Mac 电脑
  • macOS 12.3 或更高版本
  • Python 3.7 或更高版本
  • Xcode 命令行工具:xcode-select --install

三、安装步骤

您可以选择使用 Anaconda 或 pip 进行安装。请注意,使用 Apple 芯片的 Mac 和使用 Intel x86 的 Mac 的环境设置会有所不同。

使用 PyTorch 安装页面上的安装选择器,为 MPS 设备加速选择 Preview (Nightly)。MPS 后端支持是 PyTorch 1.12 正式版的一部分。PyTorch 的预览版(夜间版)将为您的设备提供最新的 MPS 支持。

安装环境

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 软件包管理器中安装。

安装 PyTorch

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

从源代码构建

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

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

四、验证安装

您可以使用以下简单的 Python 脚本来验证 MPS 是否支持:

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.")

输出结果应为:

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