Hugging Face 使用:模型库、Datasets 与 Trainer

FreeGuideOnline 最新 2026-06-16

什么是 Hugging Face?

Hugging Face 是一站式机器学习协作平台,核心围绕 Transformers 库、Datasets 库和 Trainer API 构建,让你轻松调用前沿模型、处理海量数据并完成训练。本教程从零开始,带你掌握模型库、数据集和训练流程的核心用法。

环境准备

首先确保 Python 3.8+,然后在虚拟环境中安装必要库:

pip install transformers datasets accelerate torch
  • transformers:模型加载、训练
  • datasets:数据集下载与处理
  • accelerate:多 GPU/TPU 训练后端
  • torch:PyTorch 深度学习框架(也可用 TensorFlow)

使用模型库 (Model Hub)

模型库汇集了社区分享的数千个预训练模型,涵盖 NLP、语音、视觉等任务。

探索模型库

访问 huggingface.co/models,可按任务、框架、语言等标签筛选模型。每个模型页面都提供了可直接使用的代码示例、模型卡片和文件列表。

使用 Pipeline 快速调用模型

pipeline 是最高层的抽象,自动完成文本预处理、模型推理和后处理,一行代码搞定常见任务:

from transformers import pipeline

# 情感分析
classifier = pipeline("sentiment-analysis")
result = classifier("I love Hugging Face!")
print(result)  # [{'label': 'POSITIVE', 'score': 0.999...}]

# 文本生成
generator = pipeline("text-generation", model="gpt2")
print(generator("The future of AI is", max_length=30))

你可以指定模型,例如 pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")。若无特殊需要,用默认模型即可快速验证。

使用 AutoClasses 定制化加载

当 pipeline 不够灵活时,用 AutoTokenizerAutoModel 系列手动控制每一步:

from transformers import AutoTokenizer, AutoModelForSequenceClassification

model_name = "distilbert-base-uncased-finetuned-sst-2-english"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

inputs = tokenizer("Hugging Face is amazing!", return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits
print(logits.argmax(-1))  # 预测类别索引
  • AutoTokenizer 自动匹配对应的分词器
  • AutoModelForSequenceClassification 等任务模型可自动加载权重和架构

其他常用 Auto 类:AutoModelForCausalLM(生成)、AutoModelForQuestionAnswering(问答)、AutoModelForImageClassification(图像分类)等。

使用 Datasets 库

Datasets 库让你用相同 API 访问海量公开数据集,并高效处理内存无法容纳的大数据。

加载数据集

from datasets import load_dataset

# 从 Hub 加载 IMDb 影评数据集
dataset = load_dataset("imdb")
print(dataset)
# DatasetDict({
#     train: Dataset({ features: ['text', 'label'], num_rows: 25000 })
#     test: Dataset({ features: ['text', 'label'], num_rows: 25000 })
#     unsupervised: Dataset({...})
# })

常用参数:split 选择特定分割(如 load_dataset("imdb", split="train[:1%]") 只取 1% 样本快速测试),cache_dir 指定缓存目录。

探索数据集

# 查看一条样本
print(dataset["train"][0])  # {'text': '...', 'label': 1}

# 查看列名和数据类型
print(dataset["train"].features)

# 查看数据集大小
print(len(dataset["train"]))

数据预处理

map 方法对数据集进行批量操作,并配合 tokenizer 完成文本编码:

def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)

tokenized_datasets = dataset.map(tokenize_function, batched=True)

# 移除原始文本列并设置格式
tokenized_datasets = tokenized_datasets.remove_columns(["text"])
tokenized_datasets.set_format("torch")  # 返回 PyTorch 张量
  • batched=True 可加速处理
  • truncationpadding 保证序列长度一致
  • set_format 指定返回类型("torch", "tensorflow", "numpy"

创建 DataLoader

通过 torch.utils.data.DataLoader 进一步控制批次加载:

from torch.utils.data import DataLoader

train_dataloader = DataLoader(
    tokenized_datasets["train"], shuffle=True, batch_size=8
)

使用 Trainer 进行训练

Trainer 封装了训练循环、评估、日志等,让你专注于模型和数据,无需手写繁琐代码。

准备模型

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=2
)

通常在预训练模型基础上微调,只需指定任务头对应的输出数量。

定义训练参数

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy="epoch",  # 每个 epoch 评估一次
    save_strategy="epoch",
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=3,
    weight_decay=0.01,
    logging_dir="./logs",
    logging_steps=10,
)

常用设置:learning_rate(默认 5e-5)、fp16=True(混合精度加速)、push_to_hub=True 将模型推送到 Hub。

初始化 Trainer

from transformers import Trainer

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["test"],
    tokenizer=tokenizer,
)

训练与评估

trainer.train()          # 开始训练
eval_results = trainer.evaluate()
print(eval_results)      # 返回 loss、accuracy 等指标

训练过程中会在 output_dir 自动保存检查点,你可用 trainer.save_model("./best_model") 手动保存最终模型。

快速预测

训练完成后,直接使用 pipeline 或手动推理:

from transformers import pipeline

pipe = pipeline("text-classification", model="./results/checkpoint-xxxx")
print(pipe("This movie is really great!"))

总结

通过模型库获得强力基线,Datasets 简化数据加载与预处理,Trainer 安全掌管训练全流程。这套工具链让从零到一构建高质量 NLP 模型的门槛大幅降低。熟练以上基础后,你可以尝试自定义损失、回调函数或多 GPU 分布式训练,进一步释放 Hugging Face 生态的潜力。