在M1 GPU上运行PyTorch

摘要： PyTorch团队于5月18日官宣了对M1 GPU的支持。威斯康星大学麦迪逊分校助理研究员Sebastian Raschka博士于第一时间在M1 MacBook上对新版本PyTorch做了深度学习基准测试，并将测试结果发布在个人博客上。经作者授权，本文对博客文件进行了翻译。

作者： Sebastian Raschka

授权： 已获作者授权。

正文：

今天（即2022年5月18日），PyTorch正式推出对苹果ARM M1芯片的GPU支持。对于Mac用户，这无疑是激动人心的一天。为此，当晚我就腾出时间做了实测。这篇简短的博客是我对使用M1芯片完成深度学习任务的总结，并给出了一些思考。

迄今为止我的M1体验

时间回到2021年初，我兴冲冲地卖掉了那台庞大厚重的15英寸Intel MacBook Pro，购买了便宜得多的M1 MacBook Air。时至今日，它一直表现出色，安静、轻巧、超快，电池续航时间长。

在撰写我的新书时，我注意到这台M1 MacBook不仅日常使用中感觉很快，而且对计算也做了一定的加速。例如，预处理IMDb电影数据集只需要21秒，而在我的2019 Intel MacBook Pro上曾耗时1分51秒。

同样在M1 MacBook Air上，所有使用scikit-learn工作流的运算速度也加快了！对于一些小规模PyTorch神经网络，例如我在教学中使用的各种多层感知器和卷积神经网络，都能在合理的时间内跑完。我曾对M1和GeForce 1080Ti做了LeNet-5运行时间的对比，结果运行速度近似。

尽管M1 MacBook是一台优秀的机器，但确实不适用于训练现代深度神经网络。除LeNets外，它无法去跑其它神经网络。作为一位PyTorch的早期使用者，当时我自行编译了PyTorch，但只能使用M1 CPU。

PyTorch对M1 GPU的支持

今天，PyTorch团队终于官宣了对M1 GPU的支持，对此我很兴奋。发布公报中给出的基准测试结果表明，对于训练VGG16，M1 GPU被CPU快了8倍；对于推理（评估），速度提升了约21倍。公报附则中指出测试是在配置M1 Ultra的Mac Studio上开展的，由此我认为公报中的CPU就是指M1 Ultra CPU。

那么应如何安装支持M1 GPU的PyTorch版本？我希望1.12版本中将包含对M1 GPU的支持，具体更新建议查看版本列表。现在可安装的是最新发布的nightly版本。下图来自pytorch.org的截屏：

个人推荐从终端安装，命令如下：

$ conda create -n torch-nightly python=3.8 
$ conda activate torch-nightly
$ pip install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu

之后，使用torch.device("mps")在GPU上运行PyTorch代码。这类似于对Nvidia GPU使用的torch.device("cuda")。

有意思的是，支持M1 GPU的PyTorch安装文件大小约45Mb，而支持CUDA 10.2的PyTorch安装版文件大小约750Mb。

我使用的基准测试

求知心切，我亲自动手在不同硬件上完成单回合（epoch）深度神经网络训练，分别是具有12核Intel服务器级CPU的强大深度学习工作站，以及带有M1 Pro芯片的MacBook Pro。下面给出VGG16的结果，其中CIFAR-10图像重缩放为224x224像素，这是VGG16的典型ImageNet规模：

运行细节

M1 Max和M1 Ultra结果由博客读者友情提供。此外，我给出了RTX 3060和RTX 3080 结果。

除RTX 3060之外，所有运行的batch size均为32。RTX 3060只有6 GB VRAM，能处理的batch size为8。

读者可能注意到上图中某些运行给出了两个版本，即May 18和May 22。最初的发布的May18 nightly版（torch 1.12.0.dev20220518）中存在明显的内存泄露。该问题在May 21版中得以修复。因此，我升级到May 22 nightly版，并重新运行了测试。

相比M1 Pro CPU（上数第四行），M1 Pro GPU（下数第六行）的网络训练速度提高了3.5倍。这很说明问题！

下图给出了推理性能对比。其中推理使用测试数据集评估模型：

相比M1 Pro CPU，M1 Pro GPU对推理过程的加速更为显著。例如，M1 Pro GPU的速度（1.69 分钟）是M1 Pro CPU（8.51 分钟）的5倍。

虽然相比M1 CPU，M1 GPU显著提升了神经网络的提升，但这并非颠覆性的。我们会在神经网络训练中继续考虑使用传统的GPU。当然，Mac Studio的M1 Ultra在性能上更接近于Nvidia GPU，但也应看到，这台机器在当前售价高达5000美元！

这里需要对M1 GPU的性能做一些额外说明：

• 我注意到，相比CUDA GPU，在CPU或 M1 GPU上运行卷积网络需要更多的内存，并且可能在交换（swapping）上存在问题。但我确保在MacBook Pro上训练神经网络的内存占用率从未超过 80%。

• 另一种解释是避免在M1 GPU上运行时最大化batch size（参考该建议）。然而出于公平考虑，我batch size=32运行所有的训练。其中，2080Ti和1080Ti由于11Gb VRAM的限制而无法运行。更新：我以batch size=64重复在M1 Pro GPU的运行。相比batch size=32，速度提升了约20%。

此外也应该看到，我使用了实现新特性的早期版本，该版本可能会随着时间的推移而不断改进。当然，看到新特性能运行起来，并正在积极开发中，这本身就十分令人兴奋！

如果读者也想亲自运行代码，给出我所使用脚本的Github链接。

总结

要问我是否会将配备M1 GPU的MacBook作为深度学习的首选，我当然会说不。我并不认为笔记本电脑能担当起神经网络训练的重任。原因列举如下：

• 相比工作站，笔记本的性价比不高；
• 运行时温度过高，这当然对笔记本不好；
• 在满负荷使用的情况下，笔记本就难以去完成我们日常交付的事务。

对我而言，笔记本是我的首选生产力，几乎将它用于处理所有事务，包括基本的科学运算，以及在工作站或云集群上运行之前用于调试复杂的神经网络代码。

总结我从基准测试中得到的结论：最新的PyTorch对M1 GPU的支持终于工作了，这令人兴奋！M1 GPU要比CPU工作得更顺畅和快速，但仍存在一些需要解决的问题。我期待新版本的PyTorch将会使MacBook这样的笔记本在生产效能、原型设计和深度学习调试上更具吸引力。

原文链接： Running PyTorch on the M1 GPU