Press "Enter" to skip to content

Macbook Pro M1芯片使用Pytorch进行深度学习小试

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

携手创作,共同成长!这是我参与「掘金日新计划 · 8 月更文挑战」的第30天,点击查看活动详情

 

导语

 

实验室去年为大家配新电脑,选配了一台最新的MacBook Pro,搭载了M1芯片。最近我也听说Pytorch官方宣布支持使用M1芯片的GPU加速,比CPU快了很多倍。于是便自己尝试了一下,特此记录。

 

安装Pytorch

 

最新发布的Pytorch 1.12版本已经支持对Mac的M1 GPU支持,可以直接在官网进行选择,使用相应的命令进行下载安装。

 

 

即使用如下命令进行安装

 

pip3 install torch torchvision torchaudio

 

或者conda命令为:

 

conda install pytorch torchvision torchaudio -c pytorch

 

以下是Pytorch 1.12发布时所展示的性能对比,可以看到在训练和验证普遍都快了5-20倍。

 

 

使用M1芯片进行加速

 

要想像使用服务器的GPU上进行深度学习加速,就需要将模型放到GPU上,在服务器中这个操作是通过

 

device = torch.device("cuda:0")
model = model.to(device)

 

实现,而MacBook Pro中只需要将cuda改为mps即可,即

 

device = torch.device("mps")
model = model.to(device)

 

例如,我们可以将数据和模型通过指定device的方式生成或者从cpu搬到GPU上,示例代码如下:

 

import torch
import torchvision
device = torch.device("mps")
x = torch.randn(32, 32, device=device)
model = torchvision.models.resnet18().to(device)
print(x.device)
print(next(model.parameters()).device)

 

这里的变量x直接通过指定device的方式在mps即M1芯片的GPU上生成,而模型resnet18则是从CPU生成后搬到了mps。

 

实验对比

 

看完以上示例后,我们来跑几个实验来看看效果咋样。这里使用的示例是之前博客(请参考: 简单的文本分类任务:不借助Trainer实现 )中所使用的Transformer文本分类的例子,对比的baseline也是M1芯片自己的CPU。所有的实验使用jupyterlab在一个单元格中使用魔法命令%%time统计得到运行时间。

 

实验配置

MacBook Pro(14英寸,2021年)
芯片Apple M1Pro
内存32GB

mps实验

 

%%time
model=AutoModelForSequenceClassification.from_pretrained("bert-base-uncased",num_labels=2)
optimizer = AdamW(model.parameters(), lr=2e-5)
num_epochs = 1
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler(
    name="linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps
)
device = 'mps'
model.to(device)
model.train()
for epoch in range(num_epochs):
    for i, batch in enumerate(train_dataloader):
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        progress_bar.update(1)
        
        if i>10:
            break

 

得到输出如下:

 

CPU times: user 11.2 s, sys: 8.57 s, total: 19.8 s
Wall time: 23.7 s

 

cpu实验

 

cpu实验只需要将上面的代码中的mps改为cpu即可,相关代码如下:

 

%%time
model=AutoModelForSequenceClassification.from_pretrained("bert-base-uncased",num_labels=2)
optimizer = AdamW(model.parameters(), lr=2e-5)
num_epochs = 1
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler(
    name="linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps
)
device = 'cpu'
model.to(device)
model.train()
for epoch in range(num_epochs):
    for i, batch in enumerate(train_dataloader):
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        progress_bar.update(1)
        
        if i>10:
            break

 

得到实验输出如下:

 

CPU times: user 10min 42s, sys: 11.7 s, total: 10min 53s
Wall time: 10min 56s

 

上面的实验都是训练了10个step,从以上实验可以看到,同样的一个文本分类的BERT模型,在mps上的速度是cpu上速度的二十多倍,这个速度已经基本可以满足日常小模型测试和调试的需求了,极大方便了科研学习。

 

总结

 

本文介绍了如何在MacBook Pro M1上使用Pytorch进行深度学习开发,而整个过程也非常简单,只需要将device显式的指定为’mps’。

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。