使用BERT进行fine-tuning

Github上BERT的README里面已经给出了相当详细的使用说明,GOOGLE BERT地址。
Fine-tuning就是载入预训练好的Bert模型,在自己的语料上再训练一段时间。载入模型和使用模型继续训练这部分github上代码已经帮忙做好了,我们fine-tuning需要做的工作就是在官方代码的run_classifier.py这个文件里面添加本地任务的Processor


仿照官方代码中的XnliProcessor,添加一个文本分类的fine-tuning处理类TextProcessor,它需要继承DataProcessor,实现get_train_examplesget_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文件。文件每行代表模型对每个类别预测的分数,对应顺序为TextProcessorget_labels返回的标签顺序。

你可能感兴趣的:(使用BERT进行fine-tuning)