Hugging Face 使用:模型库、Datasets 与 Trainer
什么是 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 不够灵活时,用 AutoTokenizer 和 AutoModel 系列手动控制每一步:
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可加速处理truncation和padding保证序列长度一致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 生态的潜力。