探索GitHub顶级项目：PyTorch在Apple Silicon M2平台的实战应用

引言

随着Apple Silicon芯片的普及，越来越多的开发者开始在M1/M2系列Mac上运行深度学习框架。PyTorch作为GitHub上最受欢迎的深度学习框架之一，如何在Apple Silicon M2上高效运行成为了许多开发者关心的问题。本文将带你从零开始，在M2芯片的Mac上配置PyTorch环境，并完成一个完整的图像分类实战项目。

准备工作

硬件和软件要求

Mac电脑配备Apple Silicon M2芯片
macOS Monterey (12.3+) 或更高版本
Python 3.8或更高版本
Homebrew包管理器（如果尚未安装）

为什么选择PyTorch for Apple Silicon？

PyTorch官方从1.12版本开始提供对Apple Silicon的原生支持(MPS后端)，可以充分利用M系列芯片的GPU加速能力。相比传统的CPU计算，使用MPS后端可以获得显著的性能提升。

安装步骤

1. 安装Homebrew（如未安装）

打开终端(Terminal)，执行以下命令：

代码片段

/bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"

安装完成后，将Homebrew添加到PATH环境变量：

代码片段

echo 'eval "$(/opt/homebrew/bin/brew shellenv)"' >> ~/.zshrc
source ~/.zshrc

2. 安装Python环境

推荐使用conda管理Python环境：

代码片段

brew install --cask miniforge
conda init zsh

关闭并重新打开终端后，创建新的conda环境：

代码片段

conda create -n pytorch_m2 python=3.9 -y
conda activate pytorch_m2

3. 安装PyTorch及其依赖

使用官方推荐的命令安装支持Apple Silicon的PyTorch：

代码片段

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

注意：截至本文撰写时，官方稳定版可能还不完全支持MPS后端。如果遇到问题，可以尝试nightly版本。

验证安装是否成功：

代码片段

import torch
print(torch.backends.mps.is_available())  # 应该返回True
print(torch.backends.mps.is_built())      # 应该返回True

PyTorch MPS实战：图像分类示例

下面我们通过一个完整的图像分类示例来演示如何使用PyTorch的MPS后端加速训练。

1. 准备数据集

我们将使用CIFAR-10数据集：

代码片段

import torchvision
import torchvision.transforms as transforms

# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 下载并加载训练集和测试集
trainset = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(
    trainset, batch_size=64, shuffle=True)

testset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(
    testset, batch_size=64, shuffle=False)

classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

2. 定义神经网络模型

创建一个简单的CNN模型：

代码片段

import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, len(classes))

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = torch.flatten(x, start_dim=1) # flatten all dimensions except batch
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Net()

3. MPS设备设置与模型迁移

将模型和数据移动到MPS设备：

代码片段

device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
print(f"Using device: {device}")

net.to(device) # Move model to MPS device (GPU acceleration)

4.定义损失函数和优化器

代码片段

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

5.训练网络

代码片段

for epoch in range(10): # loop over the dataset multiple times

    running_loss = 0.0

    for i, data in enumerate(trainloader):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data[0].to(device), data[1].to(device) # Move data to MPS

        # zero the parameter gradients 
        optimizer.zero_grad()

        # forward + backward + optimize 
        outputs = net(inputs) 

        loss=criterion(outputs ,labels ) 

         loss.backward() 

         optimizer.step() 

         running_loss +=loss.item() 

         if i %200 ==199: # print every200 mini-batches print('[%d,%5d] loss:%.3f'% (epoch+1,i+1,running_loss /200)) running_loss=0.

print('Finished Training')

性能对比:MPS vs CPU

为了展示MPS后端的优势，我们可以进行简单的性能对比 :

代码片段

import time 

def benchmark_device(device_type): 

     device=torch.device(device_type) 

     model=Net().to(device) 

     dummy_input=torch.randn(64 ,3 ,32 ,32 ,device=device ) 

     start_time=time.time() for _ in range (100 ): _=model(dummy_input ) end_time=time.time () return end_time -start_time 

cpu_time=benchmark_device ("cpu") mps_time=benchmark_device ("mps") 

print (f"CPU time:{cpu_time:.4f} seconds") print (f"MPS time:{mps_time:.4f} seconds") print (f"Speedup:{cpu_time/mps_time:.2f}x")

在我的MacBook Pro with M2 Max上的测试结果:

代码片段

CPU time:4 .8723 seconds MPS time :1 .3456 seconds Speedup :3 .62x

可以看到，使用MPS后端获得了约3 .6倍的加速 !

常见问题与解决方案

Q1:MPS后端不可用

如果torch.backends.mps.is_available()返回False ，请检查 :

macOS版本是否≥12 .3
PyTorch是否为≥1 .12版本
Python是否为≥3 .8

Q2:内存不足错误

由于Apple Silicon GPU与内存共享，大batch size可能导致内存不足。尝试 :

-减小batch size
-使用torch.utils.data.DataLoader的pin_memory=True参数

Q3:性能不如预期

确保你的代码确实使用了MPS设备 :

代码片段

print(next(net .parameters ()).device ) # should output "mps"

总结

本文详细介绍了如何在Apple Silicon M系列芯片上配置和使用PyTorch ，包括 :

1 .正确安装支持MPS后端的PyTorch版本
2 .将模型和数据迁移到MPS设备的方法
3 .完整的图像分类实战示例
4 .性能对比展示MPS的优势

通过利用Apple Silicon的GPU加速能力，你可以显著提升深度学习模型的训练和推理速度。随着PyTorch对M系列芯片支持的不断完善，未来还会有更好的性能表现。