我们与 Apple 的 Metal 工程团队合作，很高兴地宣布支持在 Mac 上进行 GPU 加速的 PyTorch 训练。在此之前，Mac 上的 PyTorch 训练仅利用 CPU，但随着即将发布的 PyTorch v1.12，开发人员和研究人员可以利用 Apple 芯片 GPU 大幅加快模型训练速度。这使得在 Mac 本地执行原型设计和微调等机器学习工作流程成为可能。
Metal 加速
GPU 加速训练通过 Apple 的 Metal Performance Shaders (MPS) 作为 PyTorch 的后端来实现。MPS 后端扩展了 PyTorch 框架，提供了在 Mac 上设置和运行操作的脚本和功能。MPS 通过针对每个 Metal GPU 系列的独特特性进行优化的内核来优化计算性能。新设备将机器学习计算图和原语映射到 MPS Graph 框架和 MPS 提供的优化内核上。
Apple 芯片上的训练优势
每台 Apple 芯片 Mac 都具有统一内存架构，为 GPU 提供对完整内存存储的直接访问。这使得 Mac 成为一个出色的机器学习平台，使用户能够在本地训练更大的网络或批次大小。这降低了与基于云的开发相关的成本，或减少了对额外本地 GPU 的需求。统一内存架构还降低了数据检索延迟，提高了端到端性能。
在下面的图中，您可以看到与 CPU 基准相比，GPU 加速训练和评估的性能提升
加速 GPU 训练和评估速度比仅 CPU 快（倍数）
开始使用
要开始使用，只需在运行 macOS 12.3 或更高版本并安装原生版本 (arm64) Python 的 Apple 芯片 Mac 上安装最新的 预览版 (Nightly) 即可。
您还可以在 Apple 的 Metal 页面 上了解更多关于 Metal 和 MPS 的信息。
* 测试由 Apple 于 2022 年 4 月使用配备 Apple M1 Ultra、20 核 CPU、64 核 GPU、128GB RAM 和 2TB SSD 的量产型 Mac Studio 系统进行。测试在 macOS Monterey 12.3、预发布版 PyTorch 1.12、ResNet50（批次大小=128）、HuggingFace BERT（批次大小=64）和 VGG16（批次大小=64）下进行。性能测试使用特定计算机系统进行，反映了 Mac Studio 的近似性能。