结合 Accelerator 使用 Adafactor

您已经用 Trainer 集成 Adafactor 做出了一个概念验证，使您的机器翻译模型在更低内存占用下完成训练。现在，您希望使用 Accelerator 来自定义训练循环。请构建一个使用 Adafactor 的训练循环！

compute_optimizer_size() 函数已预先定义。部分训练对象已预加载：model、train_dataloader 和 accelerator。Adafactor 已从 torch.optim 预先导入。

本练习是课程的一部分

使用 PyTorch 高效训练 AI 模型

查看课程

练习说明

在定义 optimizer 时，将模型参数传递给 Adafactor。
传入优化器状态并打印其大小。

交互式实操练习

通过完成这段示例代码来试试这个练习。

# Pass the model parameters to Adafactor
optimizer = Adafactor(params=____.____())
model, optimizer, train_dataloader = accelerator.prepare(model, optimizer, train_dataloader)

for batch in train_dataloader:
    inputs, targets = batch["input_ids"], batch["labels"]
    outputs = model(inputs, labels=targets)
    loss = outputs.loss
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()

# Pass in the optimizer state
total_size_megabytes, total_num_elements = compute_optimizer_size(optimizer.state.____())
print(f"Number of optimizer parameters: {total_num_elements:,}\nOptimizer size: {total_size_megabytes:.0f} MB")

编辑并运行代码