PyTorch 官宣适配 Apple M1，支持 GPU 加速，训练速度提升 7 倍

频道：行业资讯日期：2022-12-15 13:49:30 浏览：1185

PyTorch宣布,通过与苹果的金属工程团队合作,目前已实现在搭载苹果M1芯片的Mac上<强>使用GPU加速训练。在这之前,在Mac上进行PyTorch训练只能使用CPU,但随着PyTorch v1.12即将发布,开发者和研究者可以利用苹果M1 GPU的强大性能,从而显著提升模型训练速度。

对于Mac来说,此功能极大丰富了在Mac上执行机器学习工作流程的能力,例如在本地进行原型设计和微调。

根据PyTorch官方博客的介绍,PyTorch通过使用苹果的金属性能着色器(MPS)作为PyTorch的后端来启用GPU加速训练.MPS后端扩展了PyTorch框架,提供了在Mac上设置和运行操作的脚本和功能.MPS使用针对每个金属GPU系列的独特特性进行微调的内核来优化计算性能。新设备将机器学习计算图和原语映射到议员图框架和议员提供的调整内核上。

此外,由于搭载苹果M1芯片的Mac具有统一的内存结构,因此GPU可以直接访问完整的内存存储。这使得Mac成为机器学习的绝佳平台,使用户能够在本地训练更大的网络或批大小,并降低了与基于云算力的开发相关的成本或对额外本地GPU的需求。统一内存架构还减少了数据检索延迟,提高了端到端的性能。

下图展示了与CPU基线相比,GPU加速训练带来了显著的性能提升:

可以看的到,使用GPU可将模型训练速度提升约7倍,评估(评价)速度则最高提升约20倍。

苹果公司于2022年4月使用配备苹果M1超20核CPU, 64核GPU, 128 gb内存和2 tb SSD的Mac工作室进行了测试。系统为macOS蒙特雷12.3,安装了预发布的PyTorch 1.12,测试模型为ResNet50(批大?128),HuggingFace伯特(批大?64)和VGG16(批大?64)

虽然PyTorch已官宣适配苹果M1,不过毕竟尚未正式推出,所以并非所有基于PyTorch的模型都能使用M1的GPU加速。大家可以自行测试体验,环境要求:需要在搭载M1系列芯片的Mac上安装原生版本(arm64)的Python,以及安装最新预览版PyTorch (预览构建),并将系统升级至macOS 12.3预览版或更新的版本。