在 M1 系列上使用GPU训练Pytorch

发表于 更新于 分类于 阅读次数： 本文字数： 251 阅读时长 ≈ 1 分钟

最近 Pytorch 支持了 M1系列芯片的 GPU 接下来将演示如何在 MAC OS 上如何使用

Step1

安装 arm 版的 anconda，切记一定是 arm 版的

可以通过以下命令查看自己安装的版本是否正确

1
2
import platform
print(platform.platform())

如果是 arm 版的会输出 macOS-12.3.1-arm64-arm-64bit，如果输出带有 X86 则说明版本不对需要重新安装

Step2

安装最新预览版Pytorch

1
conda install pytorch torchvision torchaudio -c pytorch-nightly

使用 conda 有一定概率无法安装到预览版，在安装的时候可以检查 Pytorch 版本是否为 1.12.0.dev20220519

如果使用 conda 无法安装到预览版，可以使用 pip3 安装

1
pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu

Step3

到这一步就已经完成了环境的安装，可以测试是否支持 M1系列 GPU

1
2
3
import torch
a = torch.rand(5).to("mps")
# tensor([1.2810e-02, 4.8380e-01, 4.0773e-01, 2.5867e-01, 6.4366e-01], device='mps:0')

如果没有报错就说明 OK 了