由 PyTorch
我们很高兴地宣布，与苹果的 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 图框架上，以及 MPS 提供的经过调优的内核。
Apple 硅芯片上的训练优势每台 Apple 硅芯片 Mac 都具有统一的内存架构，使 GPU 可以直接访问完整的内存存储。这使得 Mac 成为机器学习的理想平台，使用户能够在本地训练更大的网络或批次大小。这降低了与基于云的开发或需要额外本地 GPU 相关的成本。统一内存架构还减少了数据检索延迟，提高了端到端性能。
在下图中，您可以看到与 CPU 基线相比，加速的 GPU 训练和评估的性能加速。
与仅 CPU 相比的加速 GPU 训练和评估加速（快多少倍）
入门要开始使用，只需在运行 macOS 12.3 或更高版本的 Apple 硅芯片 Mac 上安装最新的 预览（夜间）版本，并使用 Python 的原生版本 (arm64)。
您也可以在 Apple 的 Metal 页面 上了解