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的传入参数如下：

1
2
3
4
5
6
7
8
9
10
11
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没有提供，则使用DataCollatorWithPadding

  • train_dataset, eval_dataset为训练集，验证集

  • tokenizer为模型训练使用的tokenizer

  • model_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个参数！），我们在这里只介绍几个常见的：

1
2
3
4
5
6
7
8
9
10
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为模型评估策略
  1. “no": 不做模型评估
  2. "steps": 按训练步数（steps）进行评估，需指定步数
  3. "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。

1
2
3
4
5
6
7
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个样本。

1
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。

1
2
3
4
5
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.

1
2
from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)

设置compute_metrics函数，在评估过程中输出accuracy, f1, precision, recall四个指标。设置训练参数TrainingArguments类，设置Trainer。 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
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
)

开启模型训练。

1
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

以上为英语数据集的文本分类模型微调。

中文数据集使用sougou-mini数据集（训练集4000个样本，测试集495个样本，共5个输出类别），预训练模型采用bert-base-chinese。代码基本与英语数据集差不多，只要修改 预训练模型，数据集加载 和 最大长度为128，输出类别。以下是不同的代码之处：

1
2
3
4
5
6
7
8
9
10
11
12
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模型的动态量化）获取量化前后模型推理的各项指标。

中文数据集文本分类模型评估代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
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 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
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

1
true_labels[:10]

[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

1
pred_labels[:10]

[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

1
2
3
from sklearn.metrics import classification_report

print(classification_report(true_labels, pred_labels, digits=4))

输出结果如下：

1
2
3
4
5
6
7
8
9
10
11
              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/ ，欢迎大家访问~

欢迎关注我的公众号NLP奇幻之旅，原创技术文章第一时间推送。

欢迎关注我的知识星球“自然语言处理奇幻之旅”，笔者正在努力构建自己的技术社区。

NLP
#BERT #文本分类 #HuggingFace
NLP（六十六）使用HuggingFace中的Trainer进行BERT模型微调
https://percent4.github.io/NLP（六十六）使用HuggingFace中的Trainer进行BERT模型微调/
作者
Jclian91
发布于
2023年9月2日
许可协议