【PyTorch】使用 Mac GPUs (Apple silicon GPUs) 训练模型

根据 PyTorch 官网的文章 Introducing Accelerated PyTorch Training on Mac¹ ，从 PyTorch v1.12 release 开始支持使用 Apple silicon GPUs 加速训练模型。所以要在 Mac 上加速需要1.12或更新的版本。

文章中还给出了 M1 Ultra 在模型训练/推理时相比仅使用 CPU 训练的速度差距：

我的 Mac 芯片是 Apple M2 Pro，GPU 性能比 M1 Ultra 差一大截，但可能还是要比 CPU 快一些的。

M1 Ultra 的浮点运算速度大约是 M2 Pro 的 3 倍。

作为对比（数据来源：www.techpowerup.com）：

• RTX 3060 12GB FP16/FP32: 12.74 TFLOPS FP64: 199.0 GFLOPS
• RTX 3090 24GB FP16/FP32: 35.58 TFLOPS FP64: 556.0 GFLOPS
• RTX 4080 16GB FP16/FP32: 48.74 TFLOPS FP64: 761.5 GFLOPS
• RTX 4090 24GB FP16/FP32: 82.58 TFLOPS FP64: 1,290 GFLOPS

芯片：Apple M2 Pro
软件：conda 22.11.1, Python 3.9.6, PyTorch Version: 2.1.0.dev20230328

先查看一下 torch 能不能调用到 Apple silicon GPUs：

>>> import torch
>>> torch.backends.mps.is_available()
True
>>> torch.backends.mps.is_built()
True

可以在创建 tensor 时使用 device 参数来设置数据保存到 mps 中：

>>> import torch
>>> device = torch.device("mps")
>>> a = torch.randn((), device=device, dtype=torch.float)
>>> a
tensor(0.2838, device='mps:0')

也可以使用 .to(device) 把 CPU 上运算的 tensor 移动到 GPU 上：

>>> a = torch.tensor([1,2])
>>> a
tensor([1, 2])
>>> a.to(device)
tensor([1, 2], device='mps:0')

运行一个完整的示例²，如果没有遇到错误，就说明可以正常使用了：

import torch
import math

dtype = torch.float
device = torch.device("mps")

# Create random input and output data
x = torch.linspace(-math.pi, math.pi, 2000, device=device, dtype=dtype)
y = torch.sin(x)

# Randomly initialize weights
a = torch.randn((), device=device, dtype=dtype)
b = torch.randn((), device=device, dtype=dtype)
c = torch.randn((), device=device, dtype=dtype)
d = torch.randn((), device=device, dtype=dtype)

learning_rate = 1e-6
for t in range(2000):
    # Forward pass: compute predicted y
    y_pred = a + b * x + c * x ** 2 + d * x ** 3

    # Compute and print loss
    loss = (y_pred - y).pow(2).sum().item()
    if t % 100 == 99:
        print(t, loss)

    # Backprop to compute gradients of a, b, c, d with respect to loss
    grad_y_pred = 2.0 * (y_pred - y)
    grad_a = grad_y_pred.sum()
    grad_b = (grad_y_pred * x).sum()
    grad_c = (grad_y_pred * x ** 2).sum()
    grad_d = (grad_y_pred * x ** 3).sum()

    # Update weights using gradient descent
    a -= learning_rate * grad_a
    b -= learning_rate * grad_b
    c -= learning_rate * grad_c
    d -= learning_rate * grad_d


print(f'Result: y = {a.item()} + {b.item()} x + {c.item()} x^2 + {d.item()} x^3')

---

1. https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/ ↩︎

2. https://towardsdatascience.com/installing-pytorch-on-apple-m1-chip-with-gpu-acceleration-3351dc44d67c ↩︎