↑ 點擊藍字關(guān)注極市平臺

作者丨力元

來源丨機器之心

編輯丨極市平臺

極市導(dǎo)讀

多數(shù) PyTorch 高級庫都支持分布式訓(xùn)練和混合精度訓(xùn)練，但是它們引入的抽象化往往需要用戶學(xué)習(xí)新的 API 來定制訓(xùn)練循環(huán)。許多 PyTorch 用戶希望完全控制自己的訓(xùn)練循環(huán)，但不想編寫和維護訓(xùn)練所需的樣板代碼。Hugging Face 最近發(fā)布的新庫 Accelerate 解決了這個問題。 >>加入極市CV技術(shù)交流群，走在計算機視覺的最前沿

「Accelerate」提供了一個簡單的 API，將與多 GPU 、 TPU 、 fp16 相關(guān)的樣板代碼抽離了出來，保持其余代碼不變。PyTorch 用戶無須使用不便控制和調(diào)整的抽象類或編寫、維護樣板代碼，就可以直接上手多 GPU 或 TPU。

項目地址：https://github.com/huggingface/accelerate

通過將如下 5 行代碼添加到原始的 PyTorch 訓(xùn)練循環(huán)中，腳本即可在本地以及任何分布式設(shè)置上運行。

import torch  import torch.nn.functional as F  from datasets import load_dataset+ from accelerate import Accelerator
+ accelerator = Accelerator()- device = 'cpu'+ device = accelerator.device
  model = torch.nn.Transformer().to(device)  optim = torch.optim.Adam(model.parameters())
  dataset = load_dataset('my_dataset')  data = torch.utils.data.DataLoader(dataset, shuffle=True)
+ model, optim, data = accelerator.prepare(model, optim, data)
  model.train()  for epoch in range(10):      for source, targets in data:          source = source.to(device)          targets = targets.to(device)
          optimizer.zero_grad()
          output = model(source)          loss = F.cross_entropy(output, targets)
+         accelerator.backward(loss)-         loss.backward()
          optimizer.step()

Accelerate 甚至可以通過處理設(shè)備的放置（需要對代碼進行一些更改，但通常更安全）進一步簡化訓(xùn)練循環(huán)，代碼如下所示：

import torch  import torch.nn.functional as F  from datasets import load_dataset+ from accelerate import Accelerator
+ accelerator = Accelerator()- device = 'cpu'
+ model = torch.nn.Transformer()- model = torch.nn.Transformer().to(device)  optim = torch.optim.Adam(model.parameters())
  dataset = load_dataset('my_dataset')  data = torch.utils.data.DataLoader(dataset, shuffle=True)
+ model, optim, data = accelerator.prepare(model, optim, data)
  model.train()  for epoch in range(10):      for source, targets in data:-         source = source.to(device)-         targets = targets.to(device)
          optimizer.zero_grad()
          output = model(source)          loss = F.cross_entropy(output, targets)
+         accelerator.backward(loss)-         loss.backward()
          optimizer.step()

除了 API，Accelerate 還提供了一個 CLI 工具，方便啟動腳本之前快速配置和測試訓(xùn)練環(huán)境，然后啟動腳本。指令如下所示：

accelerate config

accelerate launch my_script.py --args_to_my_script

如果不想自己編寫訓(xùn)練循環(huán)， PyTorch 之上有許多可以替代 Accelerate 的高級庫。

Accelerate 的運作原理

accelerator = Accelerator()

除了提供要使用的主要對象之外，此行還將從環(huán)境中分析分布式訓(xùn)練運行的類型并執(zhí)行必要的初始化。用戶可以通過將 cpu = True 或 fp16 = True 傳遞給此 init 來強制進行 CPU 訓(xùn)練或混合精度訓(xùn)練。這兩個選項都可以使用腳本的啟動器進行設(shè)置。

model, optim, data = accelerator.prepare(model, optim, data)

這是 API 的主體，將準(zhǔn)備三種主要類型的對象：models (torch.nn.Module)、optimizers (torch.optim.Optimizer)、dataloaders (torch.data.dataloader.DataLoader)。

模型

模型的準(zhǔn)備包括將其包裝在適當(dāng)?shù)娜萜鳎ɡ?DistributedDataParallel）中，然后將其放置在適當(dāng)?shù)脑O(shè)備上。與普通分布式訓(xùn)練一樣，進行保存或訪問其特定的方法時，需要先通過 accelerator.unwrap_model（model）解開模型。