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