通过与 Apple Metal 工程团队的合作，我们很高兴地宣布 PyTorch 现已支持在 Mac 上进行 GPU 加速训练。此前，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 内存和 2TB SSD 的量产版 Mac Studio 系统。测试环境为 macOS Monterey 12.3、PyTorch 1.12 预发布版、ResNet50（批大小=128）、HuggingFace BERT（批大小=64）以及 VGG16（批大小=64）。性能测试是在特定的计算机系统上进行的，反映了 Mac Studio 的近似性能。