【Hugging Face】accelerate 库的常用方法和属性

Hugging Face accelerate 的常用方法和属性

accelerate 是 Hugging Face 提供的 分布式训练库，用于 加速 PyTorch 训练、自动优化多 GPU、TPU、混合精度训练，简化 数据并行、模型并行 的设置，适用于 大规模 NLP 和深度学习任务。

1. Accelerator 的常见属性

在 Accelerator 初始化后，可以访问以下常见属性：

from accelerate import Accelerator

accelerator = Accelerator()

2. Accelerator 的常用方法

3. Accelerator 详细用法

3.1. 初始化 Accelerator

from accelerate import Accelerator

# 自动检测设备（支持多 GPU/TPU）
accelerator = Accelerator()

print(f"Using device: {accelerator.device}")

3.2. 训练模型

import torch
from transformers import AutoModelForSequenceClassification, AdamW
from torch.utils.data import DataLoader, TensorDataset

# 初始化模型
model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)

# 创建数据集
dataset = TensorDataset(torch.randint(0, 30522, (100, 128)), torch.randint(0, 2, (100,)))
dataloader = DataLoader(dataset, batch_size=16, shuffle=True)

# 定义优化器
optimizer = AdamW(model.parameters(), lr=5e-5)

# 初始化 `Accelerator`
accelerator = Accelerator()

# **使用 `prepare` 自动并行化**
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

# 训练循环
for epoch in range(3):
    for batch in dataloader:
        inputs, labels = batch
        outputs = model(inputs, labels=labels)
        loss = outputs.loss

        # **自动处理混合精度**
        accelerator.backward(loss)

        optimizer.step()
        optimizer.zero_grad()

    accelerator.print(f"Epoch {epoch+1} completed.")

在 多 GPU 训练时，prepare() 自动管理设备分配，无需手动调用 .to(device)。

3.3. 评估模型

model.eval()

for batch in dataloader:
    inputs, labels = batch
    with torch.no_grad():
        outputs = model(inputs)
    
    predictions = torch.argmax(outputs.logits, dim=-1)

在 accelerate 中，无需手动管理 GPU/CPU，代码可在 单 GPU、多 GPU 或 TPU 之间无缝切换。

3.4. 训练时梯度累积

accelerator = Accelerator(gradient_accumulation_steps=4)

如果显存不足，可以使用梯度累积 gradient_accumulation_steps 减少显存占用。

3.5. 保存和加载模型

保存

accelerator.save(model.state_dict(), "model.pth")

加载

model.load_state_dict(torch.load("model.pth", map_location=accelerator.device))

3.6. 释放 GPU 内存

accelerator.free_memory()

3.7. 在 Trainer 中使用 Accelerator

如果使用 Trainer，可以启用 accelerate：

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="./results",
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    evaluation_strategy="epoch",
    fp16=True,  # **启用混合精度**
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)
trainer.train()

4. Accelerator 适用于哪些任务？

5. accelerate vs DataParallel vs DistributedDataParallel

如果你希望 无缝支持多 GPU/TPU、自动管理显存和混合精度训练，推荐 accelerate。

6. 总结

accelerate 提供 分布式训练、自动混合精度、多 GPU/TPU 支持，适用于 文本分类、机器翻译、文本生成等任务。

常见方法：

• accelerator.prepare(model, optimizer, dataloader) 自动处理设备分配
• accelerator.backward(loss) 自动管理梯度计算
• accelerator.save(model, path) 保存模型
• accelerator.free_memory() 释放显存
• accelerator.print(msg) 避免多 GPU 重复打印

如果你的 模型大、显存不足、需要多 GPU 训练，推荐使用 accelerate。