在Mac M2上安装PyTorch并启用MPS加速的详细教程与性能对比

创作时间:

作者:

@小白创作中心

在Mac M2上安装PyTorch并启用MPS加速的详细教程与性能对比

引用

CSDN

https://blog.csdn.net/weixin_71894495/article/details/144629831

本文将详细介绍如何在Mac M2上安装PyTorch并启用MPS加速，同时通过实际测试对比CPU和MPS的性能差异。

1. 安装PyTorch

在官网上可以查看安装教程，Start Locally | PyTorch。需要提前安装Python 3.9及以上版本，推荐使用Python 3.11最新版本。

使用conda安装PyTorch，在终端进入要安装的环境，执行如下命令：

conda install pytorch torchvision -c pytorch

执行完如上命令后需要等待几分钟，直到安装完毕。

2. 安装MPS

使用conda安装MPS：

conda install torch torchvision torchaudio

3. 安装测试

可以通过以下代码测试PyTorch和MPS是否安装成功：

import torch

# 查看 torch安装是否成功 并查看其版本
print(torch.__version__)

# 查看 mps是否安装成功 是否可用
print(torch.backends.mps.is_available())

# 检查 GPU 是否可用
print(torch.cuda.is_available())  # 对于 MPS，返回 False 是正常的
print(torch.backends.mps.is_available())  # 应该返回 True

# 获取 MPS 设备
mps_device = torch.device("mps")
print(mps_device)  # 输出 "mps"

如果能成功打印出PyTorch版本，并且torch.backends.mps.is_available()返回True，说明MPS可用。

4. 性能对比测试

一般来说，模型越复杂，MPS加速效果越明显。如果模型太简单，只需要几秒钟就能跑完的话，MPS加速反而不如CPU，因为MPS需要一些准备工作，把数据放入图显核心里去，如果算法太简单或者数据量太少，结果运行加速节约的时间还不如数据准备的时间长。

以下是具体的测试代码：

import torch
import torch.nn as nn
import torch.optim as optim
import time

# 设置训练参数
input_size = 4096  # 输入特征数
hidden_size = 1024  # 隐藏层神经元数
output_size = 10  # 输出类别数（例如 10 类）
num_epochs = 50  # 训练轮数
batch_size = 64  # 批量大小
learning_rate = 0.01  # 学习率

# 定义一个简单的全连接神经网络
class SimpleNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 函数：训练模型并记录训练时间
def train_model(device, num_epochs):
    # 创建数据集
    num_samples = 100000  # 数据集样本数量
    x_train = torch.randn(num_samples, input_size).to(device)
    y_train = torch.randint(0, output_size, (num_samples,)).to(device)

    # 模型、损失函数和优化器
    model = SimpleNN(input_size, hidden_size, output_size).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=learning_rate)

    # 开始计时
    start_time = time.time()

    # 训练循环
    for epoch in range(num_epochs):
        for i in range(0, num_samples, batch_size):
            # 获取当前批量数据
            inputs = x_train[i:i+batch_size]
            labels = y_train[i:i+batch_size]

            # 前向传播
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            # 反向传播和优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

    # 结束计时
    end_time = time.time()

    # 返回训练时间
    return end_time - start_time

# 主程序
if __name__ == "__main__":
    # 设备列表
    devices = {
        "CPU": torch.device("cpu"),
        "MPS": torch.device("mps") if torch.backends.mps.is_available() else None,
    }

    # 分别测试 CPU 和 MPS
    results = {}
    for device_name, device in devices.items():
        if device is None:
            print(f"\nSkipping {device_name} as it is not available.")
            continue
        print(f"\nTraining on {device_name}...")
        training_time = train_model(device, num_epochs)
        results[device_name] = training_time
        print(f"Training time on {device_name}: {training_time:.2f} seconds")

    # 打印对比结果
    print("\n--- Training Time Comparison ---")
    for device_name, time_taken in results.items():
        print(f"{device_name}: {time_taken:.2f} seconds")

测试环境为Mac Mini M2（8+10）16G+1T。