M1芯片Pytorch环境实现GPU加速

1 基础需求

• 一台 M1 系列芯片的 Mac 设备（系统为 Monterey 及以上）；
• arm64 的 Python（建议 3.8 版本及以上）；
• 1.12 版本及以上的 PyTorch

2 安装 PyTorch

官网链接

conda install pytorch torchvision torchaudio -c pytorch

3.1 测试使用GPU加速前

device = torch.device("cpu")

结果：15-35s运行一组

epoch 1, train_loss 0.991261,test_loss 1.073757,train_acc 0.574238,test_acc 0.573529,Time used 23.796787s
epoch 2, train_loss 0.929583,test_loss 1.055528,train_acc 0.574238,test_acc 0.552941,Time used 21.342122s
……
epoch 60, train_loss 0.521239,test_loss 0.510754,train_acc 0.817109,test_acc 0.823529,Time used 22.526722s

3.2 使用GPU加速后效果

注意：驱动写mps，不是gpu

device = torch.device("mps")

结果：3-9s运行一组

epoch 1, train_loss 1.063146,test_loss 1.121217,train_acc 0.565388,test_acc 0.161765,Time used 9.087449s
epoch 2, train_loss 1.049286,test_loss 1.121504,train_acc 0.582104,test_acc 0.161765,Time used 3.308240s
epoch 3, train_loss 1.037148,test_loss 1.114957,train_acc 0.578171,test_acc 0.364706,Time used 3.299558s
……
epoch 198, train_loss 0.666738,test_loss 0.650006,train_acc 0.761062,test_acc 0.744118,Time used 4.019118s
epoch 199, train_loss 0.663524,test_loss 0.661568,train_acc 0.761062,test_acc 0.741176,Time used 4.174166s
epoch 200, train_loss 0.664021,test_loss 0.638504,train_acc 0.758112,test_acc 0.741176,Time used 4.890875s

如果想简单测试一下，可以参考此链接
开心✌️
如果是25s一组，epoch=200，则需要1.4小时
如果是5s一组，epoch=200，则需要27分钟
如果无其他占用情况，和网上的测试差不多6倍的提速，是可以达到的，不错😁

4 参考

配置参考博客