携手创作，共同生长！这是我参与「日新方案 8 月更文应战」的第30天，点击查看活动详情

导语

试验室去年为大家配新电脑，选配了一台最新的MacBook Pro，搭载了M1芯片。最近我也听说Pytorch官方宣告支撑运用M1芯片的GPU加快，比CPU快了很多倍。所以便自己尝试了一下，特此记载。

装置Pytorch

最新发布的Pytorch 1.12版别现已支撑对Mac的M1 GPU支撑，能够直接在官网进行选择，运用相应的指令进行下载装置。

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

即运用如下指令进行装置

pip3 install torch torchvision torchaudio

或许conda指令为：

conda install pytorch torchvision torchaudio -c pytorch

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

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

运用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’。