PyTorch实现苹果M1芯片GPU加速：训练速度提升7倍，性能最高提升21倍

5月18日，PyTorch官网宣布，通过与Apple的Metal工程团队合作，目前已可以支持在搭载M1芯片的Mac上使用GPU加速PyTorch训练。而在此之前，在Mac上进行PyTorch训练仅能利用CPU。

但随着PyTorch v1.12版本的发布，开发和研究人员可以利用Apple Silicon CPU的优势，大大加快模型训练速度。这解锁了在Mac上执行机器学习工作流程的能力，例如在本地进行原型设计和微调。

据PyTorch官网介绍，其使用Apple的Metal Performance Shaders（MPS）作为PyTorch的后端来加速GPU训练。MPS后端扩展了PyTorch框架，提供了在Mac上设置和运行操作的脚本和功能。MPS通过针对每个Metal GPU系列的独特特性进行微调的内核来优化计算性能。新设备将机器学习计算图和原语映射到MPS提供的MPS Graph框架和优化内核上。

此外，因为每台搭载Apple Silicon的Mac都有统一的内存架构，为GPU提供了直接访问完整内存存储的能力。这让Mac成为机器学习的绝佳平台，也使用户能够在本地训练更大的网络或批大小。同时，这也降低了与基于云的开发相关的成本或对额外本地GPU的需求。此外，统一内存架构还减少了数据检索延迟，提高了端到端的性能。

在下图中，显示了与CPU基线相比，加速GPU训练和评估所带来的性能提升：

图源： PyTorch官网

可以看出，使用GPU可将模型训练速度提升约7倍，评估（Evaluation）速度最高可提升约21倍。

以上是Apple于2022年4月使用配备Apple M1 Ultra、20核CPU、64核GPU、128GB内存和2TB SSD的Mac Studio系统进行的测试结果。系统为macOS Monterey 12.3、预发布的PyTorch 1.12，测试模型为ResNet50（batch size=128）、HuggingFace BERT（batch size=64）和VGG16（batch size=64）。性能测试是使用特定的计算机系统进行的，反映了Mac Studio的大致性能。

对于想要体验的用户，以下为具体要求：只需在搭载M1芯片的Mac上安装原生版本（arm64）的Python，并将系统升级至macOS 12.3或更高版本，然后去官网下载最新的PyTorch预览版就可以。

参考链接：

1.Introducing Accelerated PyTorch Training on Mac | PyTorch

2.Start Locally | PyTorch