预训练模型给各类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再训练好的语言模型的使用方法