Hugging Face 的 Trainer类用法

珊珊而川

已于 2025-01-15 16:28:00 修改

阅读量427

点赞数 6

分类专栏：研0人工智能每天一个专业知识文章标签：机器学习深度学习人工智能

于 2025-01-15 16:10:17 首次发布

本文链接：https://blog.csdn.net/weixin_63681863/article/details/145162615

版权

研0人工智能每天一个专业知识专栏收录该内容

51 篇文章

订阅专栏

一、使用方法

Hugging Face 的 Trainer 类是一个高级API，用于简化训练、评估和预测的流程。以下是如何使用 Trainer 类的基本步骤：

1. 导入必要的类和函数

首先，您需要导入 Trainer 类以及其他可能需要的类或函数。

from transformers import Trainer, TrainingArguments

2. 准备数据集

您需要准备一个训练数据集和一个验证数据集。数据集可以是 Dataset 类的实例，通常使用 load_dataset 函数来加载。

from datasets import load_dataset
raw_datasets = load_dataset("your_dataset_name")
train_dataset = raw_datasets["train"]
eval_dataset = raw_datasets["test"]

3. 初始化模型和分词器

根据您的任务选择一个预训练模型和相应的分词器。

from transformers import AutoTokenizer, AutoModelForSequenceClassification
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=num_labels)

4. 准备训练参数

定义训练参数，如学习率、训练批次大小、评估频率等。

training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=64,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir="./logs",
    logging_steps=10,
)

5. 创建 `Trainer` 实例

使用上述准备好的参数创建 Trainer 实例。

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,  # 这是一个自定义函数，用于计算评估指标
)

6. 训练模型

调用 train 方法开始训练模型。

trainer.train()

7. 评估模型

在训练期间或之后，可以使用 evaluate 方法来评估模型。

eval_result = trainer.evaluate()

8. 预测

使用 predict 方法进行预测。

predictions = trainer.predict(test_dataset)

9. 保存和加载模型

保存和加载模型以便后续使用。

trainer.save_model("./save_directory")
model = AutoModelForSequenceClassification.from_pretrained("./save_directory")

注意事项

Trainer 类提供了许多参数和回调函数，可以根据需要进行自定义。
compute_metrics 函数必须返回一个字典，其中包含评估指标。
数据集需要是 Dataset 类的实例，并且数据集中的数据需要是模型可以理解的格式。
Trainer 类也支持多GPU和TPU训练。
使用 Trainer 类可以极大地简化训练流程，特别是对于常见的NLP任务，如文本分类、命名实体识别等。

二、使用了Transformers提供的Trainer类，不能说他使用了Transformers模型是吗？

是的，使用 Trainer 类并不一定意味着使用了 Transformers 提供的预训练模型。

Trainer 是 Hugging Face Transformers 框架中的通用训练工具，它确实是为 Transformers 的预训练模型设计的，但理论上，它也可以用于非 Transformers 的模型，只要这些模型的输入和输出格式符合 Trainer 的要求。
例如，如果你自己定义了一个 PyTorch 模型，且其输入格式类似于 Transformers 模型（例如，包含 input_ids, attention_mask, labels 等），你也可以使用 Trainer 来训练这个模型。

但在代码1中，明确使用了 AutoModelForSequenceClassification.from_pretrained("roberta-base")，所以可以肯定代码1中使用了 Hugging Face 提供的 Transformers 预训练模型。

三、非`Transformers` 的预训练模型也可以使用Trainer类吗？

是的，Trainer 类可以用于训练非 Transformers 的模型，只要满足以下条件：

模型格式：模型是一个标准的 PyTorch nn.Module，并且支持前向传播返回损失值或 logits。
数据格式：数据集包含 input_ids, attention_mask, 和 labels（或者可以自定义数据格式）。
自定义训练逻辑（如果需要）：通过修改 Trainer 的默认行为，可以为任意模型定义特定的训练逻辑。

例子：使用 `Trainer` 训练一个自定义模型

你可以这样使用 Trainer 训练一个普通的 PyTorch 模型：

from torch import nn
from transformers import Trainer, TrainingArguments

# 自定义 PyTorch 模型
class CustomModel(nn.Module):
    def __init__(self, input_dim, num_labels):
        super().__init__()
        self.fc = nn.Linear(input_dim, num_labels)

    def forward(self, input_ids=None, labels=None):
        logits = self.fc(input_ids)
        loss = None
        if labels is not None:
            loss_fn = nn.CrossEntropyLoss()
            loss = loss_fn(logits, labels)
        return {"loss": loss, "logits": logits}

# 创建自定义模型
model = CustomModel(input_dim=768, num_labels=2)

# 定义数据集 (确保与 Trainer 的输入兼容)
from datasets import Dataset
import numpy as np

data = {"input_ids": np.random.rand(100, 768).astype(np.float32), "labels": np.random.randint(0, 2, 100)}
dataset = Dataset.from_dict(data)

# 设置训练参数
training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=8,
    evaluation_strategy="epoch",
)

# 创建 Trainer 实例
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    eval_dataset=dataset,
)

# 开始训练
trainer.train()

总结：

Trainer 不局限于 Transformers 的预训练模型，可以用于任何符合要求的 PyTorch 模型。
如果你在代码中明确使用了 Transformers 的模型加载器（如 AutoModel*），那就可以说是使用了 Transformers 的预训练模型。

四、Trainer类封装这个用法可以自动进行分布式训练吗？我的设备有2块GPU，这个训练过程竟然把2块GPU都使用了？

是的，Trainer类可以自动支持分布式训练，并且它会根据设备环境（如多GPU）自动启用适合的分布式训练模式。这也是Trainer类强大的一个特点。

为什么`Trainer`会使用两块GPU？

Trainer基于Hugging Face的transformers库，它底层依赖PyTorch，并且与torch.distributed深度集成。
如果检测到有多个GPU可用，它会自动启用DataParallel或DistributedDataParallel模式（视具体配置而定），以充分利用所有的可用GPU资源进行训练。

分布式训练的具体机制：

默认行为：
- 当调用Trainer时，系统会自动检查设备环境（通过torch.cuda.device_count()等方法）。
- 如果有多个GPU，Trainer会自动启用DataParallel或DistributedDataParallel进行训练。
分布式通信后端：
- Trainer使用torch.distributed的通信后端（如NCCL）来协调多个GPU之间的参数同步。
- 每个GPU处理一部分数据（称为“mini-batch”），然后在每次迭代后通过通信后端同步参数。
训练参数控制：
- 通过TrainingArguments中的参数，可以显式控制分布式训练的行为。
  - device：Trainer会自动分配到cuda设备。
  - per_device_train_batch_size：指定单个GPU的batch size，多个GPU时会自动按设备数扩展。
  - distributed：内部自动检测，不需要用户手动设置。

如何确认`Trainer`是否使用了两块GPU？

查看日志输出：
- 如果Trainer检测到多个GPU，会在训练日志中显示类似信息：
```
Using distributed mode with NCCL backend
```
- 还可以看到每个设备的负载情况和训练进度。
监控GPU资源：
- 使用nvidia-smi命令查看GPU的利用情况。如果两块GPU的显存和计算资源都有明显使用，说明分布式训练正在运行。

如何手动控制分布式训练行为？

如果需要更明确地控制分布式训练，可以通过调整TrainingArguments参数实现：

示例：

training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=16,  # 单GPU的batch size
    per_device_eval_batch_size=16,
    num_train_epochs=1,
    weight_decay=0.01,
    logging_dir="./logs",
    logging_steps=10,
    save_total_limit=2,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    greater_is_better=True,
    report_to="none",
    fp16=True,  # 启用混合精度训练（更快）
    dataloader_num_workers=4,  # 加速数据加载
    local_rank=-1,  # -1表示不强制指定GPU分布，自动分配
)