以往，我们在使用HuggingFace在训练BERT模型时，代码写得比较复杂，涉及到数据处理、token编码、模型编码、模型训练等步骤，从事NLP领域的人都有这种切身感受。事实上，HugggingFace中提供了datasets模块（数据处理）和Trainer函数，使得我们的模型训练较为方便。关于datasets模块，可参考文章NLP（六十二）HuggingFace中的Datasets使用。
本文将会介绍如何使用HuggingFace中的Trainer对BERT模型微调。
Trainer
Trainer是HuggingFace中的模型训练函数，其网址为：https://huggingface.co/docs/transformers/main_classes/trainer 。
Trainer的传入参数如下：
model: typing.Union[transformers.modeling_utils.PreTrainedModel, torch.nn.modules.module.Module] = None
args: TrainingArguments = None
data_collator: typing.Optional[DataCollator] = None
train_dataset: typing.Optional[torch.utils.data.dataset.Dataset] = None
eval_dataset: typing.Union[torch.utils.data.dataset.Dataset, typing.Dict[str, torch.utils.data.dataset.Dataset], NoneType] = None
tokenizer: typing.Optional[transformers.tokenization_utils_base.PreTrainedTokenizerBase] = None
model_init: typing.Union[typing.Callable[[], transformers.modeling_utils.PreTrainedModel], NoneType] = None
compute_metrics: typing.Union[typing.Callable[[transformers.trainer_utils.EvalPrediction], typing.Dict], NoneType] = None
callbacks: typing.Optional[typing.List[transformers.trainer_callback.TrainerCallback]] = None
optimizers: typing.Tuple[torch.optim.optimizer.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None)
preprocess_logits_for_metrics: typing.Union[typing.Callable[[torch.Tensor, torch.Tensor], torch.Tensor], NoneType] = None )
参数解释：
model为预训练模型args为TrainingArguments（训练参数）类data_collator会将数据集中的元素组成一个batch，默认使用default_data_collator()，如果tokenizer没有提供，则使用DataCollatorWithPaddingtrain_dataset,eval_dataset为训练集，验证集tokenizer为模型训练使用的tokenizermodel_init为模型初始化compute_metrics为验证集的评估指标计算函数callbacks为训练过程中的callback列表optimizers为模型训练中的优化器preprocess_logits_for_metrics为模型评估阶段前对logits的预处理
TrainingArguments为训练参数类，其网址为：https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments，传入参数非常多（transformers版本4.32.1中有98个参数！），我们在这里只介绍几个常见的：
output_dir: stroverwrite_output_dir: bool = False
evaluation_strategy: typing.Union[transformers.trainer_utils.IntervalStrategy, str] = 'no'
per_gpu_train_batch_size: typing.Optional[int] = None
per_gpu_eval_batch_size: typing.Optional[int] = None
learning_rate: float = 5e-05
num_train_epochs: float = 3.0
logging_dir: typing.Optional[str] = None
logging_strategy: typing.Union[transformers.trainer_utils.IntervalStrategy, str] = 'steps'
save_strategy: typing.Union[transformers.trainer_utils.IntervalStrategy, str] = 'steps'save_steps: float = 500
report_to: typing.Optional[typing.List[str]] = None
参数解释：
output_dir为模型输出目录evaluation_strategy为模型评估策略
- “no": 不做模型评估
- “steps”: 按训练步数（steps）进行评估，需指定步数
- “epoch”: 每个epoch训练完后进行评估
per_gpu_train_batch_size,per_gpu_eval_batch_size为每个GPU上训练集和测试集的batch size，也有CPU上的对应参数learning_rate为学习率logging_dir为日志输出目录logging_strategy为日志输出策略，同样有no, steps, epoch三种，意义同上save_strategy为模型保存策略，同样有no, steps, epoch三种，意义同上report_to为模型训练、评估中的重要指标（如loss, accurace）输出之处，可选择azure_ml, clearml, codecarbon, comet_ml, dagshub, flyte, mlflow, neptune, tensorboard, wandb，使用all会输出到所有的地方，使用no则不会输出。
下面我们使用Trainer进行BERT模型微调，给出英语、中文数据集上文本分类的示例代码。
BERT微调
使用datasets模块导入imdb数据集（英语影评数据集，常用于文本分类），加载预训练模型bert-base-cased的tokenizer。
import numpy as np
from transformers import AutoTokenizer, DataCollatorWithPadding
import datasets
checkpoint = 'bert-base-cased'
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
raw_datasets = datasets.load_dataset('imdb')
查看数据集，有train（训练集）、test（测试集）、unsupervised（非监督）三部分，我们这里使用训练集和测试集，各自有25000个样本。
raw_datasets
DatasetDict({
train: Dataset({
features: ['text', 'label'],
num_rows: 25000
})
test: Dataset({
features: ['text', 'label'],
num_rows: 25000
})
unsupervised: Dataset({
features: ['text', 'label'],
num_rows: 50000
})
})
创建数据tokenize函数，对文本进行tokenize，最大长度设置为300，同时使用data_collector为DataCollatorWithPadding。
def tokenize_function(sample):
return tokenizer(sample['text'], max_length=300, truncation=True)
tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
加载分类模型，输出类别为2.
from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)
设置compute_metrics函数，在评估过程中输出accuracy, f1, precision, recall四个指标。设置训练参数TrainingArguments类，设置Trainer。
from transformers import Trainer, TrainingArguments
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
def compute_metrics(pred):
labels = pred.label_ids
preds = pred.predictions.argmax(-1)
precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='weighted')
acc = accuracy_score(labels, preds)
return {
'accuracy': acc,
'f1': f1,
'precision': precision,
'recall': recall
}
training_args = TrainingArguments(output_dir='imdb_test_trainer', # 指定输出文件夹，没有会自动创建
evaluation_strategy="epoch",
per_device_train_batch_size=32,
per_device_eval_batch_size=32,
learning_rate=5e-5,
num_train_epochs=3,
warmup_ratio=0.2,
logging_dir='./imdb_train_logs',
logging_strategy="epoch",
save_strategy="epoch",
report_to="tensorboard")
trainer = Trainer(
model,
training_args,
train_dataset=tokenized_datasets["train"],
eval_dataset=tokenized_datasets["test"],
data_collator=data_collator, # 在定义了tokenizer之后，其实这里的data_collator就不用再写了，会自动根据tokenizer创建
tokenizer=tokenizer,
compute_metrics=compute_metrics
)
开启模型训练。
trainer.train()
| Epoch | Training Loss | Validation Loss | Accuracy | F1 | Precision | Recall |
|---|---|---|---|---|---|---|
| 1 | 0.364300 | 0.223223 | 0.910600 | 0.910509 | 0.912276 | 0.910600 |
| 2 | 0.164800 | 0.204420 | 0.923960 | 0.923941 | 0.924375 | 0.923960 |
| 3 | 0.071000 | 0.241350 | 0.925520 | 0.925510 | 0.925759 | 0.925520 |
TrainOutput(global_step=588, training_loss=0.20003824169132986, metrics={'train_runtime': 1539.8692, 'train_samples_per_second': 48.705, 'train_steps_per_second': 0.382, 'total_flos': 1.156249755e+16, 'train_loss': 0.20003824169132986, 'epoch': 3.0})
以上为英语数据集的文本分类模型微调。
中文数据集使用sougou-mini数据集（训练集4000个样本，测试集495个样本，共5个输出类别），预训练模型采用bert-base-chinese。代码基本与英语数据集差不多，只要修改 预训练模型，数据集加载 和 最大长度为128，输出类别。以下是不同的代码之处：
import numpy as np
from transformers import AutoTokenizer, DataCollatorWithPadding
import datasets
checkpoint = 'bert-base-chinese'
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
data_files = {
"train": "./data/sougou/train.csv", "test": "./data/sougou/test.csv"}
raw_datasets = datasets.load_dataset("csv", data_files=data_files, delimiter=",")
...
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=5)
...
输出结果如下：
| Epoch | Training Loss | Validation Loss | Accuracy | F1 | Precision | Recall |
|---|---|---|---|---|---|---|
| 1 | 0.849200 | 0.115189 | 0.969697 | 0.969449 | 0.970073 | 0.969697 |
| 2 | 0.106900 | 0.093987 | 0.973737 | 0.973770 | 0.975372 | 0.973737 |
| 3 | 0.047800 | 0.078861 | 0.973737 | 0.973740 | 0.974117 | 0.973737 |
模型评估
在上述模型评估过程中，已经有了模型评估的各项指标。
本文也给出单独做模型评估的代码，方便后续对模型做量化时（后续介绍BERT模型的动态量化）获取量化前后模型推理的各项指标。
中文数据集文本分类模型评估代码如下：
import torch
from transformers import AutoModelForSequenceClassification
MAX_LENGTH = 128
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
checkpoint = f"./sougou_test_trainer_{
MAX_LENGTH}/checkpoint-96"
model = AutoModelForSequenceClassification.from_pretrained(checkpoint).to(device)
from transformers import AutoTokenizer, DataCollatorWithPadding
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
import pandas as pd
test_df = pd.read_csv("./data/sougou/test.csv")
test_df.head()
| text | label | |
|---|---|---|
| 0 | 届数比赛时间比赛地点参加国家和地区冠军亚军决赛成绩第一届1956-1957英国11美国丹麦6... | 0 |
| 1 | 商品属性材质软橡胶带加浮雕工艺+合金彩色队徽吊牌规格162mm数量这一系列产品不限量发行图案... | 0 |
| 2 | 今天下午，沈阳金德和长春亚泰队将在五里河相遇。在这两支球队中沈阳籍球员居多，因此这场比赛实际... | 0 |
| 3 | 本报讯中国足协准备好了与特鲁西埃谈判的合同文本，也在北京给他预订好了房间，但特鲁西埃爽约了！... | 0 |
| 4 | 网友点击发表评论祝贺中国队夺得五连冠搜狐体育讯北京时间5月6日，2006年尤伯杯羽毛球赛在日... | 0 |
import numpy as np
import time
s_time = time.time()
true_labels, pred_labels = [], []
for i, row in test_df.iterrows():
row_s_time = time.time()
true_labels.append(row["label"])
encoded_text = tokenizer(row['text'], max_length=MAX_LENGTH, truncation=True, padding=True, return_tensors='pt').to(device)
# print(encoded_text)
logits = model(**encoded_text)
label_id = np.argmax(logits[0].detach().cpu().numpy(), axis=1)[0]
pred_labels.append(label_id)
if i % 100 == 0:
print(i, (time.time() - row_s_time)*1000, label_id)
print("avg time: ", (time.time() - s_time) * 1000 / test_df.shape[0])
0 229.3872833251953 0
100 362.0314598083496 1
200 311.16747856140137 2
300 324.13792610168457 3
400 406.9099426269531 4
avg time: 352.44047810332944
true_labels[:10]
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
pred_labels[:10]
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
from sklearn.metrics import classification_report
print(classification_report(true_labels, pred_labels, digits=4))
precision recall f1-score support
0 0.9900 1.0000 0.9950 99
1 0.9691 0.9495 0.9592 99
2 0.9900 1.0000 0.9950 99
3 0.9320 0.9697 0.9505 99
4 0.9895 0.9495 0.9691 99
accuracy 0.9737 495
macro avg 0.9741 0.9737 0.9737 495
weighted avg 0.9741 0.9737 0.9737 495
总结
本文介绍了如何使用HuggingFace中的Trainer对BERT模型微调。可以看到，使用Trainer进行模型微调，代码较为简洁，且支持功能丰富，是理想的模型训练方式。
本文项目代码已开源至Github，网址为：https://github.com/percent4/PyTorch_Learning/tree/master/huggingface_learning 。
本人已开通个人博客网站，网址为：https://percent4.github.io/ ，欢迎大家访问~