携手创作,共同成长!这是我参与「掘金日新计划 · 8 月更文挑战」的第30天, 简单的文本分类任务:不借助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’。