Github上BERT的README里面已经给出了相当详细的使用说明,GOOGLE BERT地址。
Fine-tuning就是载入预训练好的Bert模型,在自己的语料上再训练一段时间。载入模型和使用模型继续训练这部分github上代码已经帮忙做好了,我们fine-tuning需要做的工作就是在官方代码的run_classifier.py这个文件里面添加本地任务的Processor
。
仿照官方代码中的XnliProcessor
,添加一个文本分类的fine-tuning处理类TextProcessor
,它需要继承DataProcessor
,实现get_train_examples
,get_labels
方法。
使用的文本数据不需要分词(bert代码会帮你做掉),以'\t'
为分隔分为两列。第一列是文本的标签,第二列是输入文本。
train.tsv:
2 比奔驰宝马比不了比相应的的车辆小五你是第一
0 美国车就是费油既然你买的起就能烧得起油
class TextProcessor(DataProcessor):
"""Processor for the Text classification"""
# 读入训练集语料 (.tsv即使用'\t'分隔的csv文件)
def get_train_examples(self, data_dir):
return self._create_examples(
self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")
# 读入evalation使用的语料, 如果--do_eval=False, 可以不实现
def get_dev_examples(self, data_dir):
return self._create_examples(
self._read_tsv(os.path.join(data_dir, "dev.tsv")),
"dev_matched")
# 读入测试集数据, --do_predict的时候使用
def get_test_examples(self, data_dir):
return self._create_examples(
self._read_tsv(os.path.join(data_dir, "test.tsv")), "test")
# 文本分类的标签
def get_labels(self):
return ["0", "1", "2"]
def _create_examples(self, lines, set_type):
"""Creates examples for the training and dev sets."""
examples = []
# lines是_read_tsv()方法读入的数据, 每一行样本按照"\t"切成一个list
for (i, line) in enumerate(lines):
# guid是样本的id, 保证唯一性就可以了
guid = "%s-%s" % (set_type, i)
# 输入的文本数据 (不需要分词)
text_a = tokenization.convert_to_unicode(line[1])
# 只有一个输出文本, text_b置为None
text_b = None
if set_type == "test":
# 这边假设test数据没有标签, 随便给个标签就行了
label = "0"
else:
# 训练和验证使用的数据标签
label = tokenization.convert_to_unicode(line[0])
examples.append(
InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
return examples
还需要在def main()
函数中添加任务
processors = {
"cola": ColaProcessor,
"mnli": MnliProcessor,
"mrpc": MrpcProcessor,
"xnli": XnliProcessor,
"text": TextProcessor
}
最后在shell脚本中添加以下内容,运行:
# 下载的预训练模型文件的目录
export BERT_BASE_DIR=/path/to/bert/uncased_L-12_H-768_A-12
# fine-tuning使用的文件目录 (目录包含train.tsv, dev.tsv, test.tsv文件)
export TEXT_DIR=/path/to/text
CUDA_VISIBLE_DEVICES=0 python run_classifier.py \
--task_name=text \
--do_train=true \
--do_eval=true \
--data_dir=$TEXT_DIR \
--vocab_file=$BERT_BASE_DIR/vocab.txt \
--bert_config_file=$BERT_BASE_DIR/bert_config.json \
--init_checkpoint=$BERT_BASE_DIR/bert_model.ckpt \
--max_seq_length=128 \
--train_batch_size=32 \
--learning_rate=2e-5 \
--num_train_epochs=3.0 \
--output_dir=/tmp/text_output/
run_classifier.py
使用的参数含义在文件开头都有解释,这里就不再赘述了。
如果想利用fine-tuning好的模型来对test数据进行预测,可以参考以下shell脚本
export BERT_BASE_DIR=/path/to/bert/uncased_L-12_H-768_A-12
export DATA_DIR=/path/to/text
# 前面fine-tuning模型的输出目录
export TRAINED_CLASSIFIER=/tmp/text_output/
python run_classifier.py \
--task_name=text \
--do_predict=true \
--data_dir=$DATA_DIR \
--vocab_file=$BERT_BASE_DIR/vocab.txt \
--bert_config_file=$BERT_BASE_DIR/bert_config.json \
--init_checkpoint=$TRAINED_CLASSIFIER \
--max_seq_length=128 \
--output_dir=/tmp/text_output/pred/
预测完成后会在/tmp/text_output/pred/目录下生成一个test_results.tsv文件。文件每行代表模型对每个类别预测的分数,对应顺序为TextProcessor
中get_labels
返回的标签顺序。