打开网易新闻 查看更多图片

整理 | 于轩 责编 | 张红月

出品 | CSDN(ID:CSDNnews)

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或更高版本,然后去官网(https://pytorch.org/get-started/locally)下载最新的PyTorch预览版就可以。

  • https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/

  • https://pytorch.org/get-started/locally/

成就一亿技术人