初试在M1 pro芯片跑一下pytorch的MPS

由于《深度学习入门2:自制框架》，从step 52开始，需要学会使用GPU替代CPU进行计算，手上只有一台macbook pro m1的本子，配置如下：

所以，尝试用pytorch的MPS来替代书中的CuPy部分代码，这一块在CSDN的很多博主文章里都有过说明了，我就是记录在此加以备忘。
PyTorch MPS (Multi-Process Service)是 PyTorch 中的一种分布式训练方式。它是基于Apple的MPS(Metal Performance Shaders) 框架开发的。MPS可以在多核的苹果设备上加速tensor的运算。MPS使用了多个设备上的多个核心来加速模型的训练。它可以将模型的计算过程分配到多个核心上，并且可以在多个设备上进行训练，从而提高训练速度。
测试代码如下：

import torch

# this ensures that the current macOS version is at least 12.3+
print(torch.backends.mps.is_available())
# this ensures that the current PyTorch installation was built with MPS activated.
print(torch.backends.mps.is_built())

N = 1000000000

device = torch.device("mps")

cpu_a = torch.randn([1, N])
cpu_b = torch.randn([N, 1])
# print(N, cpu_a.device, cpu_b.device)

gpu_a = torch.randn([1, N], device=device)
gpu_b = torch.randn([N, 1], device=device)


# print(N, gpu_a.device, gpu_b.device)


def cpu_run():
    c = torch.matmul(cpu_a, cpu_b)
    return c


def gpu_run():
    c = torch.matmul(gpu_a, gpu_b)
    return c


import timeit

cpu_time = timeit.timeit(cpu_run, number=3)
gpu_time = timeit.timeit(gpu_run, number=3)
print('warmup', cpu_time, gpu_time)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

跑下来的时间情况是这样的：

warmup 30.063021500012837 0.9484448339790106
1

看得出来，差异显著，后面就开始慢慢磨，把代码给实现出来。至少走到目前，书中前面所有的步骤，代码经过调整都能跑通，就边修改代码结构，边继续学习后面的步骤。