【AI大模型】Transformers大模型库（十一）：Trainer训练类

LDG_AGI

发布于 2024-08-13 16:18:18

1.8K0

一、引言

这里的Transformers指的是huggingface开发的大模型库，为huggingface上数以万计的预训练大模型提供预测、训练等服务。

🤗 Transformers 提供了数以千计的预训练模型，支持 100 多种语言的文本分类、信息抽取、问答、摘要、翻译、文本生成。它的宗旨是让最先进的 NLP 技术人人易用。 🤗 Transformers 提供了便于快速下载和使用的API，让你可以把预训练模型用在给定文本、在你的数据集上微调然后通过 model hub 与社区共享。同时，每个定义的 Python 模块均完全独立，方便修改和快速研究实验。 🤗 Transformers 支持三个最热门的深度学习库： Jax, PyTorch 以及 TensorFlow — 并与之无缝整合。你可以直接使用一个框架训练你的模型然后用另一个加载和推理。

本文重点介绍Trainer训练类。

二、Trainer训练类

2.1 概述

2.2 使用示例

from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments
from datasets import load_dataset

# 1. 加载数据集
# 假设我们使用的是Hugging Face的内置数据集，例如SST-2
dataset = load_dataset('sst2')  # 或者使用你自己的数据集

# 2. 数据预处理，可能需要根据模型进行Tokenization
# 以BERT为例，使用AutoTokenizer
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
def tokenize_function(examples):
    return tokenizer(examples["sentence"], truncation=True)
tokenized_datasets = dataset.map(tokenize_function, batched=True)

# 3. 准备训练参数
training_args = TrainingArguments(
    output_dir='./results',          # 输出目录
    num_train_epochs=3,              # 总的训练轮数
    per_device_train_batch_size=16,  # 每个GPU的训练批次大小
    per_device_eval_batch_size=64,   # 每个GPU的评估批次大小
    warmup_steps=500,                # 预热步数
    weight_decay=0.01,               # 权重衰减
    logging_dir='./logs',            # 日志目录
    logging_steps=10,
)

# 4. 准备模型
model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased")

# 5. 创建Trainer并开始训练
trainer = Trainer(
    model=model,                         # 要训练的模型
    args=training_args,                  # 训练参数
    train_dataset=tokenized_datasets['train'],  # 训练数据集
    eval_dataset=tokenized_datasets['validation'], # 验证数据集
)

# 开始训练
trainer.train()

整个流程是机器学习项目中的标准流程：数据准备、模型选择、参数设置、训练与评估。每个步骤都是为了确保模型能够高效、正确地训练，以解决特定的机器学习任务：

加载数据集 (load_dataset('sst2'))：这行代码是使用Hugging Face的datasets库加载SST-2数据集，这是一个情感分析任务的数据集。如果你使用自定义数据集，需要相应地处理和加载数据。
数据预处理 (tokenizer(examples["sentence"], truncation=True))：在训练模型之前，需要将文本数据转换为模型可以理解的格式。这里使用AutoTokenizer对文本进行分词（Tokenization），truncation=True意味着如果句子超过模型的最大输入长度，将截断超出部分。这一步是将文本转换为模型输入的张量格式。
训练参数 (TrainingArguments)：这部分定义了训练过程的配置，包括训练轮数(num_train_epochs)、每个设备的训练和评估批次大小、预热步数(warmup_steps)、权重衰减(weight_decay)等。这些参数对训练效率和模型性能有重要影响。
准备模型 (AutoModelForSequenceClassification.from_pretrained())：这里选择或初始化模型，AutoModelForSequenceClassification是用于序列分类任务的模型，from_pretrained方法加载预训练的模型权重。选择的模型（如BERT的“bert-base-uncased”）是基于任务需求的。
创建Trainer (Trainer)：Trainer是Transformers库中的核心类，它负责模型的训练和评估流程。它接收模型、训练参数、训练数据集和评估数据集作为输入。Trainer自动处理了训练循环、损失计算、优化器更新、评估、日志记录等复杂操作，使得训练过程更加简洁和高效。
开始训练 (trainer.train())：调用此方法开始模型的训练过程。Trainer会根据之前设定的参数和数据进行模型训练，并在每个指定的步骤打印日志，训练完成后，模型的权重会保存到指定的输出目录。