PyTorch宣布支持苹果M1芯片GPU加速：训练快6倍，推理提升21倍

2022-05-19 16:18:32

机器之心报道

编辑：泽南、蛋酱

对于 Mac 用户来说，这是令人激动的一天。

今年 3 月，苹果发布了其自研 M1 芯片的最终型号 M1 Ultra，它由 1140 亿个晶体管组成，是有史以来个人计算机中最大的数字。苹果宣称只需 1/3 的功耗，M1 Ultra 就可以实现比桌面级 GPU RTX 3090 更高的性能。

随着用户数量的增长，人们已经逐渐接受使用 M1 芯片的计算机，但作为一款 Arm 架构芯片，还有人在担心部分任务的兼容性问题。

昨天，通过与苹果 Metal 团队工程师合作，PyTorch 官方宣布已正式支持在 M1 版本的 Mac 上进行 GPU 加速的 PyTorch 机器学习模型训练。

此前，Mac 上的 PyTorch 训练仅能利用 CPU，但随着即将发布的 PyTorch v1.12 版本，开发和研究人员可以利用苹果 GPU 大幅度加快模型训练。现在，人们可以在 Mac 上相对高效地执行机器学习工作，例如在本地进行原型设计和微调。

苹果芯片的 AI 训练优势

PyTorch GPU 训练加速是使用苹果 Metal Performance Shaders (MPS) 作为后端来实现的。MPS 后端扩展了 PyTorch 框架，提供了在 Mac 上设置和运行操作的脚本和功能。MPS 使用针对每个 Metal GPU 系列的独特特性进行微调的内核能力来优化计算性能。新设备将机器学习计算图和原语映射到 MPS Graph 框架和 MPS 提供的调整内核上。

每台搭载苹果自研芯片的 Mac 都有着统一的内存架构，让 GPU 可以直接访问完整的内存存储。PyTorch 官方表示，这使得 Mac 成为机器学习的绝佳..，让用户能够在本地训练更大的网络或批大小。

这降低了与基于云算力的开发相关的成本或对额外的本地 GPU 算力需求。统一内存架构还减少了数据检索延迟，提高了端到端性能。

可以看到，与 CPU 基线相比，GPU 加速实现了成倍的训练性能提升：

上图是苹果于 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 的大致性能。

有开发者推测，鉴于谷歌云服务中使用的英伟达 T4 在 FP32 任务上的浮点性能为 8 TFLOPS，而 M1 Ultra 的图形计算能力大概在 20 TFLOPS 左右。在最有利情况下，可以期望的 M1 Ultra 速度提升或可达到 2.5 倍。

若想使用最新的加速能力，你需要在使用 M1 系列芯片的 Mac 电脑上安装原生版本（arm64）的 Python，并将系统升级至 macOS 12.3 预览版或更新的版本。

开发者亲测：加速效果显著

虽然官方已宣布提供支持，但目前还不是所有在 PyTorch 上的模型都能用 M1 芯片集成的 GPU 加速，你也可以花几分钟进行一下测试。

机器学习研究者，捷克理工大学博士 Dmytro Mishkin 对多个模型的推理进行了测试，结果显示，大多数图像分类架构都提供了很好的加速。对于一些自定义代码（比如 kornia），可能无法正常工作。

各个测试结果如下：

首先是经典的卷积神经网络 VGG16，从 2.23 秒提升到 0.5 秒：

接下来是大部分芯片发布会上都会跑的 Resnet50，它在 M1 GPU 上的速度较慢，不升反降，从 0.549 秒到 0.592 秒：

但 ResNet18 的提速惊人，从 0.243 秒到 0.024 秒：

AlexNet 的速度对比为 0.126 秒 vs0.005 秒，速度提升了几十倍：

尝试一下视觉 transformer 模型，在 M1 CPU 上的速度是 1.855 秒，在 M1 GPU 上则运行崩溃了……

EfficientNetB0 实现了 2.5 倍的加速：

EfficientNetB4 实现了 3.5 倍加速：

ConvMixer 运行良好，从 11 秒提速到 2.8 秒：

Dmytro Mishkin 也表示，使用 M1 芯片集成的 GPU 加速只需要预热一下模型，没有同步命令。和 CUDA 不同，无需异步执行。

威斯康星大学麦迪逊分校助理教授 Sebastian Raschka 也对 M1 芯片的 GPU 机器学习能力进行了一番测试，他使用的芯片是 M1 和 M1 Pro。

看上去，M1 CPU 似乎比 M1 GPU 更快。但 LeNet-5 是一个非常小的网络，而 MNIST 是一个非常小的数据集。如果用 rescaled CIFAR-10 图像再试一次，结果如下：

与 M1 Pro CPU（正数第二行）和 M1 Pro GPU（倒数第二行）相比，M1 Pro GPU 训练网络的速度提高了一倍。

可见，M1 系列芯片的 GPU 加速结果非常可观，在部分情况下已能满足开发者的需求。不过我们知道在 M1 Ultra 这样的芯片中也有 32 核的神经网络引擎，目前却只有苹果自己的 Core ML 框架支持使用该部分获得加速。

不知启用了 Neural Engine 之后，M1 芯片的 AI 推理速度还能提升多少？

参考内容：

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

http://sebastianraschka.com/blog/2022/pytorch-m1-gpu.html

5 月 21 日，机器之心组织的「ACL 2022 线上论文分享会」将在机器之心机动组视频号直播，整体日程如下，欢迎大家扫码关注。

转载请联系本公众号获得授权

投稿或寻求报道：content@jiqizhixin.com

上一篇：什么是超表面光学技术？
下一篇：给全球 10 多万员工大幅涨薪，微软是怎么想的？

小编推荐

1 开场白格式以及内容（开场白怎么写范文）
大家好，小美今天来为大家解答开场白格式以及内容以下问题，开场白怎么写范文很多人还不知道，现在让我们一起来看看吧！1、开场白的答题格式
2 关于组织参观2024第二十四届中国国际环卫与市政设施及清洗设备展览会等活动的通知
存眷我们龙年大吉不忘初心为实现清洗保洁强省方针而奋斗总运营、编纂/郑亚森迎接到场河南省清洗保洁行业协会 1、阅读并赞成协会章程，
3 穿破石——一味治肿瘤的良药
导读穿破石为桑科柘属植物构棘的根。其根系繁密蓬勃，生于山坡灌丛或疏林、乱石堆中。本咀嚼微甘、微吃力，性平，功能止咳化痰、祛风利湿、
4 厨房里的8大危险行为，千万不要做！
厨房是家里最受迎接的处所，每到饭点，香气四溢。然而厨房里也隐藏着多种危险，若是处理欠妥，或者造成严重危险。这8种危险行为万万不要做！
5 泰国国王去世（泰国王死亡）
大家好，小伟今天来为大家解答泰国国王去世以下问题，泰国王死亡很多人还不知道，现在让我们一起来看看吧！1、有影响，有一些景点会被关闭，
6 征稿通知
第十届尼山世界文明论坛将于2024年7月10日至11日在山东曲阜举办。为深入进修贯彻习近平总书记对妇女儿童和妇联工作的主要论说精神，弘扬中华民
7 沙嗲牛肉（沙嗲牛肉面）
大家好，小豪今天来为大家解答沙嗲牛肉以下问题，沙嗲牛肉面很多人还不知道，现在让我们一起来看看吧！1、您想问的是沙爹牛肉还是沙嗲牛肉好
8 乳腺癌多是拖出来的，有这5大症状，趁早就医！
★中医进修（按摩艾灸刮痧）视频合集乳腺癌的发病率始终居高不下，稳居女性恶性肿瘤发病率第一名。或者会激发乳腺癌的身分有好多，包罗喝

PyTorch宣布支持苹果M1芯片GPU加速：训练快6倍，推理提升21倍

最新文章

小编推荐