使用 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()

将梯度清零
之前已把数据移动到设备
移除手动移动数据的代码
执行前向传播
计算交叉熵损失与梯度
将 loss.backward 替换为 accelerator
更新模型参数与学习率

变更摘要

使用 Accelerator 之前

需手动将数据移到设备
- inputs.to(device)
- targets.to(device)
用 loss.backward() 计算梯度

使用 Accelerator 之后

自动设备放置与数据并行
- accelerator.prepare(model)
- accelerator.prepare(dataloader)
用 accelerator.backward(loss) 同步梯度
自定义循环
友好、与硬件无关、可扩展、易维护

Passons à la pratique !

使用 PyTorch 高效训练 AI 模型