用自己的数据增量训练预训练语言模型

预训练模型给各类NLP任务的性能带来了巨大的提升，预训练模型通常是在通用领域的大规模文本上进行训练的。而很多场景下，使用预训练语言模型的下游任务是某些特定场景，如金融，法律等。这是如果可以用这些垂直领域的语料继续训练原始的预训练模型，对于下游任务往往会有更大的提升。

以BERT为例，利用huggingface的tranformers介绍一下再训练的方式：

1. 定义tokenizer

bert的预训练模式一般分为，Masked language model (MLM)与 next sentence prediction(NSP)，主要利用MLM在自己的语料上进行预训练

from transformers import RobertaConfig,BertTokenizer
from transformers import BertForMaskedLM as Model
from transformers import MaskedLMDataset,Split
from transformers import DataCollatorForLanguageModeling
from transformers.trainer_utils import get_last_checkpoint

# 定义tokenizer
tokenizer = BertTokenizer.from_pretrained(retrained_bert_path, max_len=max_seq_length)

2. 定义预训练模型的参数

# 定义预训练模型的参数
config = RobertaConfig(
    vocab_size=tokenizer.vocab_size,
    max_position_embeddings=max_seq_length,
    num_attention_heads=12,
    num_hidden_layers=12,
    type_vocab_size=2,
)

预训练的模式为MLM，直接调用 DataCollatorForLanguageModeling API即可方便得以自己的语料定义生成器。

retrained_model = Model(config=config)
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm=True, mlm_probability=0.15)

3. 加载MLM训练数据

train_data = MaskedLMDataset(data_file=train_file,
                                 tokenizer=tokenizer,
                                 tag=train_tags,
                                 max_seq_length=max_seq_length,
                                 mode=Split.train,
                                 overwrite_cache=overwrite_cache)
train_data = [feature.convert_feature_to_dict() for feature in train_data]

 

4. 开始预训练

这里可以设置的参数有，输入端的batch_size、语料文件、tokenizer，训练过程方面则有 训练轮数epochs、batch_size 以及保存频率。经过这些简单的即可成功训练好一个基于MLM的bert模型了（损失loss降到0.5左右就可以了），也可以通过MLM模型所带的接口来做MLM预测，当然我们这里需要的只是bert的权重。

设置训练参数

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
        output_dir=out_model_path,
        overwrite_output_dir=True,
        num_train_epochs=train_epoches,
        per_device_train_batch_size=batch_size,
        save_steps=2000,
        save_total_limit=2,
        prediction_loss_only=True,
    )

训练

trainer = Trainer(
        model=retrained_model,
        args=training_args,
        train_dataset=train_data,
        data_collator=data_collator,
    )

last_checkpoint = get_last_checkpoint(training_args.output_dir)
if last_checkpoint is not None:
    train_result = trainer.train(resume_from_checkpoint=last_checkpoint)
else:
    train_result = trainer.train()

保存模型

trainer.save_model()  # Saves the tokenizer too for easy upload

本篇介绍了 增量训练预训练语言模型的方法，下一篇将介绍fine-tunning再训练好的语言模型的使用方法