用 Accelerator 訓練模型

使用 PyTorch 高效訓練 AI 模型

Dennis Lee

Data Engineer, Amazon

Trainer 與 Accelerator

比較 Accelerator 與 Trainer 在易用性與客製化能力的圖表。

自訂訓練迴圈

Trainer：沒有自訂訓練迴圈
生成式 AI 的進階任務常需兩個網路

生成對抗網路的示意圖。

¹ https://www.aitude.com/basics-of-generative-adversarial-network-model/

Trainer 與 Accelerator

比較 Accelerator 與 Trainer 在易用性與客製化能力的圖表。

修改基本訓練迴圈

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)

    outputs = model(inputs)

    loss = outputs.loss

    loss.backward()

    optimizer.step()
    scheduler.step()

將梯度歸零
將資料移到指定裝置：.to(device)
執行前向傳播
計算交叉熵損失
反向傳播計算梯度
更新模型參數與學習率

建立 Accelerator 物件

Accelerator 提供分散式訓練介面

from accelerate import Accelerator
accelerator = Accelerator(
    device_placement=True
)

device_placement（bool，預設 True）：預設自動處理裝置配置

定義模型與最佳化器

載入預訓練模型

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-cased", return_dict=True)

以 Adam 最佳化模型參數

from torch.optim import Adam

optimizer = Adam(params=model.parameters(), lr=2e-5)

定義學習率排程器

from transformers import get_linear_schedule_with_warmup

lr_scheduler = get_linear_schedule_with_warmup(
      optimizer=optimizer,

      num_warmup_steps=num_warmup_steps,

      num_training_steps=num_training_steps)

optimizer（obj）：PyTorch 最佳化器，如 Adam
num_warmup_steps（int）：線性增加 lr 的步數，設為 int(num_training_steps * 0.1)
num_training_steps（int）：總訓練步數，設為 len(train_dataloader) * num_epochs

準備模型以高效訓練

prepare 方法會處理裝置配置

model, optimizer, dataloader, lr_scheduler = \
    accelerator.prepare(model,

                        optimizer,

                        dataloader,

                        lr_scheduler)

用 Accelerator 建立訓練迴圈

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)

將梯度歸零
先前已將資料移至裝置

用 Accelerator 建立訓練迴圈

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

將梯度歸零
先前已將資料移至裝置
移除手動移動資料的程式碼

用 Accelerator 建立訓練迴圈

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)

    loss = outputs.loss
    loss.backward()

將梯度歸零
先前已將資料移至裝置
移除手動移動資料的程式碼
執行前向傳播
計算交叉熵損失與梯度

用 Accelerator 建立訓練迴圈

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)
    loss = outputs.loss

    accelerator.backward(loss)

    optimizer.step()
    scheduler.step()

將梯度歸零
先前已將資料移至裝置
移除手動移動資料的程式碼
執行前向傳播
計算交叉熵損失與梯度
以 accelerator 取代 loss.backward
更新模型參數與學習率

變更摘要

使用 Accelerator 之前

需要手動把資料移到裝置
- inputs.to(device)
- targets.to(device)
用 loss.backward() 計算梯度

使用 Accelerator 之後

自動裝置配置與資料平行化
- accelerator.prepare(model)
- accelerator.prepare(dataloader)
用 accelerator.backward(loss) 處理梯度同步
迴圈可自訂
介面友善、硬體無關、可擴充、易維護

一起來練習吧！

使用 PyTorch 高效訓練 AI 模型