初试在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)

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

warmup 30.063021500012837 0.9484448339790106

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

  • 19
    点赞
  • 9
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值