使用transformers库微调Bert进行文本分类

最近尝试利用HuggingFace🤗的transformers库在pytorch下进行Bert文本分类的微调，找了很多中文blog，主要是对数据的处理这块没有比较详细的说明，不知道怎么处理dataset的格式，因此在这里做一下记录。

依赖包

pytorch
transformers
scikit-learn

预训练模型加载

预训练模型加载这块HuggingFace在transformers库中封装得非常好，没什么太多要讲的：

args.pretrain 这里填写模型名称或者你自己准备好的预训练模型，各种预训练模型可以从https://huggingface.co/models查找和下载，需要包含三个文件（config.json、vocab.txt、pytorch_model.bin）。此处以Bert为例，准备bert-base-chinese模型，args.pretrain 为存放模型三个文件的路径。

由于下游任务是文本分类任务，因此model使用transformer.BertForSequenceClassification，也可以根据需要选择其他模型。

from from transformers import BertForSequenceClassification, BertTokenizerFast

tokenizer = BertTokenizerFast.from_pretrained(args.pretrain)
model = BertForSequenceClassification.from_pretrained(args.pretrain, num_labels=2, output_hidden_states=False)

模型微调

模型微调这里使用transformers封装好的Trainer模块，参数含义基本上都比较一目了然，这里设置了早停、根据precision加载最佳模型。值得注意的是在模型保存时会保存多个checkpoint，因此evaluation_strategy、save_total_limit要设置一下以免保存过程中爆硬盘，Bert一个checkpoint保存下来差不多要1GB……

from transformers import Trainer, TrainingArguments, EarlyStoppingCallback
from sklearn.metrics import classification_report, precision_score, \
    recall_score, f1_score, 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='binary')
    acc = accuracy_score(labels, preds)
    return {
        'accuracy': acc,
        'f1': f1,
        'precision': precision,
        'recall': recall
    }

training_args = TrainingArguments(
        output_dir=args.save_path,  # 存储结果文件的目录
        overwrite_output_dir=True,
        num_train_epochs=args.epoch,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size,
        learning_rate=1e-5,
        eval_steps=500,
        load_best_model_at_end=True,
        metric_for_best_model="precision",  # 最后载入最优模型的评判标准，这里选用precision最高的那个模型参数
        weight_decay=0.01,
        warmup_steps=500,
        evaluation_strategy="steps",  # 这里设置每100个batch做一次评估，也可以为“epoch”，也就是每个epoch进行一次
        logging_strategy="steps",
        save_strategy='steps',
        logging_steps=100,
        save_total_limit=3,
        seed=2021,
        logging_dir=args.logging_dir  # 存储logs的目录
    )

trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_set,
        eval_dataset=valid_set,
        tokenizer=tokenizer,
        compute_metrics=compute_metrics,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],  # 早停Callback
    )

数据预处理

终于要说一下数据预处理阶段了，为什么最先进行的操作要放在最后讲呢？因为读者朋友可能也发现了，transformers🤗封装得非常好，以至于整个pipeline中需要进行自定义的就是数据预处理这块。其实也非常简单，你需要在创建 Trainer的时候传入train_dataset和 eval_dataset，这两个数据集的类型都是torch.utils.data.Dataset，PyTorch的 Dataset 处理详见另一篇文章。那么这里需要对 __getitem__方法进行一些修改，使其返回一个dict，里面有包含Bert输入所需的元素。

Talk is cheap，this is the code:

from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __init__(self, file, tokenizer, max_len=512):
        assert os.path.exists(file)
        data = open(file, 'r', encoding='utf-8').read().strip().split('\n')
        texts = [x.split('\t')[0][:max_len-2] for x in data]
        labels = [int(x.split('\t')[1]) for x in data]
        self.encodings = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
        self.labels = torch.tensor(labels)

    def __getitem__(self, idx):
        item = {key: val[idx] for key, val in self.encodings.items()}
        item['labels'] = self.labels[idx]
        return item

    def __len__(self):
        return len(self.labels)

假设数据格式为“text\tlabel”，即文本和标签中间用制表符\t隔开，每行一条数据，形如：

我的梦想是星辰大海 1

雄心万丈躺在床上 0

在 __init__初始化方法中读入数据，之后截断至最大设置长度max_len（由于tokenizer会自动补上[CLS]和[SEP]，因此这里需要对最大长度做-2处理。对字符串分割出文本与标签后，使用tokenizer(texts, padding=True, truncation=True, return_tensors='pt')得到Bert所需要的输入encoding（transformers.BatchEncoding对象，encoding.data包含了'input_ids'，'token_type_ids'，'attention_mask'三个张量对象，通常情况下在微调任务中不需要进行额外处理），之后将标签转为张量对象就基本完成了数据集初始化流程。

比较关键的是__getitem__方法的处理，这里需要返回一个dict对象，里面需要包含input_ids, token_tpye_ids, attention_mask, labels四个key（以Bert为例，其他模型可能会有稍许不同，注意这里虽然是单条数据，但是标签的key称为“labels”），然后返回该dict对象即可（即代码示例中的item。

开始训练！

最后一步，开始训练！静静等待模型训练完即可。

trainer.train()

代码示例

还没有push到GitHub上去，后续改。

Reference

Fine-tuning pretrained NLP models with Huggingface’s Trainer