NLP（六十六）使用HuggingFace中的Trainer进行BERT模型微调
以往，我们在使用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 datasetscheckpoint = '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_supportdef compute_metrics(pred):labels = pred.label_idspreds = 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 datasetscheckpoint = '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 AutoModelForSequenceClassificationMAX_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, DataCollatorWithPaddingtokenizer = AutoTokenizer.from_pretrained(checkpoint)import pandas as pdtest_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 times_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_reportprint(classification_report(true_labels, pred_labels, digits=4))
precision recall f1-score support0 0.9900 1.0000 0.9950 991 0.9691 0.9495 0.9592 992 0.9900 1.0000 0.9950 993 0.9320 0.9697 0.9505 994 0.9895 0.9495 0.9691 99accuracy 0.9737 495macro 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/ ，欢迎大家访问~