colab上基于tensorflow2的BERT中文文本多分类finetuning

整体背景

本文实现了在colab环境下基于tf-nightly-gpu的BERT中文多分类，如果你在现阶段有实现类似的功能的需求，详细这篇文章会给你带来一些帮助。

准备工作

1.环境：

硬件环境：

直接使用谷歌提供的免费训练环境colab，选择GPU

软件环境：

tensorflow：tensorflow2.1.0版本对BERT的支持有些问题，现象是可以训练但预测时无法正常加载模型（稍后代码里会详述），因此改为选择tf-nightly（主线版本随时会有新变更，如果担心有影响可以选择打tag的dev版本，如笔者验证2.2.0.dev20200218是可以的）

预训练模型：基于tf2.x的中文预训练模型（https://tfhub.dev/tensorflow/bert_zh_L-12_H-768_A-12/1）【注：谷歌基于tf1.x版本公布过一个中文预训练模型（BERT-base，Chinese），根据注释其在tf2.0环境下应该无法使用】

分类代码逻辑：谷歌已经帮我们准备好了基于BERT的分类主程序，其不随tensorflow一起提供，而是单独放在了tensorflow/models下面。models目前也分为发布版和主线版本，本文为了保证分类代码逻辑的稳定，选择发布版models v2.1.0【注：models主线版本近期（2020年2月）有频繁变更，可以看到的是正在对BERT分类逻辑做小规模的重构优化，所以当models 2.2.0的时候，可能分类逻辑会有较大变化】

适配代码：谷歌提供的代码肯定不能直接满足我们自己的需求，所以要对代码进行修改调整，本文也提供了基于models2.1.0版本修改出来的BERT分类适配代码

2.训练数据：

本文从百度百科上爬取了10w左右的词条，按照人物(human)、自然(nature)、地点(poi)、组织(org)、生活(life)等等标注成大概十几个类别。把每个百科词条中的信息按照标题、概述、正文、infobox等分别提取出来并进行简单的清洗操作后，直接连接到一起形成一个长文本。再与百度百科id以及分类标签按照"id\tinfo\tlabel"的格式组织成3列，形成类似于这种格式的训练数据：

bkid    info    label

---------------

1         xxxx    life

2        xxxx    human

……

实际数据看起来是这样的：

 数据准备好后重新洗牌并按照70%:20%:10%的比例拆分成训练集、验证集和测试集。

代码适配

bert相关代码都在tensorflow/models下面，将models 2.1.0压缩包下载下来后后，BERT分类代码位于：models2.1.0/official/nlp/bert目录下

$ tree .
.
|-- bert_cloud_tpu.md
|-- classifier_data_lib.py # 分类数据方法库
|-- common_flags.py # 通用命令行参数
|-- create_finetuning_data.py # 生成tfrecord格式的微调数据（依赖classifier_data_lib.py）
|-- create_pretraining_data.py # 生成预训练数据（只有重新预训练整个模型时才使用，本次不用）
|-- custom_metrics.py # 自定义文件，用于计算自定义metrics
|-- do_pred_data.py # 自定义文件，用于生成预测数据
|-- do_predict.py # 自定义文件，用于模型训练完成后执行预测
|-- export_tfhub.py
|-- export_tfhub_test.py
|-- __init__.py
|-- input_pipeline.py # 用于加载tfrecord格式的训练数据
|-- model_saving_utils.py # 用于模型保存
|-- README.md
|-- run_classifier.py # finetune分类主程序
|-- run_pretraining.py # 预训练主程序（本次不用）
|-- run_squad.py # squad主程序（本次不用）
|-- squad_lib.py
|-- squad_lib_sp.py
|-- tf1_checkpoint_converter_lib.py
|-- tf2_albert_encoder_checkpoint_converter.py
|-- tf2_encoder_checkpoint_converter.py
|-- tokenization.py # 分词器
|-- tokenization_test.py
`-- utils.py #自定义文件，用于绘制自定义metrics图表

0 directories, 25 files

主要修改点有二：

1. 添加自定义的数据处理器及处理逻辑

2. 原finetune代码只有训练集和验证集的代码流程，本文在整个训练结束后添加了测试集的相关流程。

3. 原代码只系统了训练逻辑，本文添加了训练完成之后的预测逻辑（暂未包括生产环境部署）。

下面对主要修改点进行说明：

1.数据格式转换逻辑：

tensorflow推荐使用tfrecord格式的数据，因此将数据集转换成tfrecord格式并保存下来，便于后续重复使用

文件1：create_finetuning_data.py 

a. 添加新的数据处理器：

 1 def generate_classifier_dataset():
 2   """Generates classifier dataset and returns input meta data."""
 3   assert FLAGS.input_data_dir and FLAGS.classification_task_name
 4 
 5   processors = { 
 6       "cola": classifier_data_lib.ColaProcessor,
 7       "mnli": classifier_data_lib.MnliProcessor,
 8       "mrpc": classifier_data_lib.MrpcProcessor,
 9       "qnli": classifier_data_lib.QnliProcessor,
10       "sst-2": classifier_data_lib.SstProcessor,
11       "xnli": classifier_data_lib.XnliProcessor,
12       "bdbk": classifier_data_lib.BdbkProcessor, # 添加新的数据处理器
13   }
14 ......

 

文件2：classifier_data_lib.py

a. 添加新Processor的处理逻辑：

谷歌预设的数据处理器以及后续的分类逻辑中，都只使用了训练集和验证集，本文在此基础上，增加了测试集和预测集的数据处理和使用

 1 class BdbkProcessor(DataProcessor):
 2     """Processor for bdbk data set."""
 3 
 4     def get_train_examples(self, data_dir):
 5         """See base class."""
 6         return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")
 7 
 8     def get_dev_examples(self, data_dir):
 9         """See base class."""
10         return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")
11 
12     def get_test_examples(self, data_dir):
13         """See base class."""
14         return self._create_examples(self._read_tsv(os.path.join(data_dir, "test.tsv")), "test")
15 
16     def get_predict_examples(self, data_dir):
17         return self._create_examples(self._read_tsv(os.path.join(data_dir, "predict.tsv")), "predict")
18 
19     def get_labels(self):
20         """See base class."""
21         return [
22             'human',
23             'poi',
24             'nature',
25             ... 此处根据你的需要，定义自己的类别标签
26         ]
27     @staticmethod
28     def get_processor_name():
29         """See base class."""
30         return "BDBK"
31 
32     def _create_examples(self, lines, set_type):
33         """Creates examples for the training and dev sets."""
34         examples = []
35         for (i, line) in enumerate(lines):
36             guid = "%s-%s" % (set_type, i)
37             if i == 0: # 如果你的数据有标题行，就加上这两句代码跳过
38                 continue
39             text_a = self.process_text_fn(line[1]) # line[1]列是文本
40             label = self.process_text_fn(line[2]) if len(line) >= 3 else "na" # line[2]列是标签
41             examples.append(InputExample(guid=guid, text_a=text_a, text_b=None, label=label))
42         return examples

 

2.训练、验证及测试逻辑：

对tfrecord数据读取、模型训练结果存储等功能进行调整

文件1：run_classifier.py

run_classifier是分类主程序入口，涵盖了预训练模型加载、分类模型构建及编译、输入数据读取、模型训练和评估、产出训练结果等全过程。

其中分类模型构建环节中没有添加其他分类网络，即直接用BERT预训练模型的输出对接dropout（抑制过拟合）层后，直接接fc层产出各个类别的概率。此处也可以自定义其他网络结构，比如BERT的输出对接一个TextCNN，然后再接fc层（这种网络结构也可以理解成BERT作为word embedding层，TextCNN作为分类模型）。本文中采用第一种默认结构，即BERT+dropout+fc，网络模型结构如下：

a. 添加一些epoch结尾的评估metrics，便于评估训练效果

 1     eval_data_list = list(evaluation_dataset.as_numpy_iterator())
 2     custom_metric = custom_metrics.Metrics(labels_list, valid_data=eval_data_list[0])
 3 
 4     if custom_callbacks is not None:
 5       custom_callbacks += [custom_metric, summary_callback, checkpoint_callback]
 6     else:
 7       custom_callbacks = [custom_metric, summary_callback, checkpoint_callback]
 8 
 9     history = bert_model.fit(
10         x=training_dataset,
11         validation_data=evaluation_dataset,
12         steps_per_epoch=steps_per_epoch,
13         epochs=epochs,
14         validation_steps=eval_steps,
15         callbacks=custom_callbacks)
16 
17     return bert_model, history, custom_metric

文件2：model_saving_utils.py

a. 在export_bert_model()最后添加训练结果保存逻辑

    # tf.train.Checkpoint API was used via custom training loop logic.
    else:
      checkpoint = tf.train.Checkpoint(model=model)

      # Restores the model from latest checkpoint.
      latest_checkpoint_file = tf.train.latest_checkpoint(checkpoint_dir)
      assert latest_checkpoint_file
      logging.info('Checkpoint file %s found and restoring from '
                   'checkpoint', latest_checkpoint_file)
      checkpoint.restore(
          latest_checkpoint_file).assert_existing_objects_matched()

  model.save(model_export_path, include_optimizer=True, overwrite=True, save_format='h5') # 添加训练结果保存
  return

b. 新增custom_metrics.py，用于在on_epoch_end中计算自定义metrics

from sklearn.metrics import f1_score, recall_score, precision_score, classification_report
from tensorflow.keras.callbacks import Callback
import numpy as np

class Metrics(Callback):
    def __init__(self, labels_list, valid_data):
        super(Metrics, self).__init__()
        self.validation_data = valid_data
        self.val_f1s = []
        self.val_recalls = []
        self.val_precisions = []
        self.reports = []
        self.labels_list = labels_list

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        val_predict = np.argmax(self.model.predict(self.validation_data[0]), -1) 
        val_targ = self.validation_data[1]
        if len(val_targ.shape) == 2 and val_targ.shape[1] != 1:
            val_targ = np.argmax(val_targ, -1) 

        # 分别计算macro f1, recall, precision
        _val_precision = precision_score(val_targ, val_predict, average='macro')
        self.val_precisions.append(_val_precision)
        logs['val_precision'] = _val_precision


        _val_recall = recall_score(val_targ, val_predict, average='macro')
        self.val_recalls.append(_val_recall)
        logs['val_recall'] = _val_recall

        _val_f1 = f1_score(val_targ, val_predict, average='macro')
        self.val_f1s.append(_val_f1)
        logs['val_f1'] = _val_f1
        remain_list = []
        val_union = np.union1d(val_targ, val_predict)
        for i, v in enumerate(self.labels_list):
            if i in val_union:
                remain_list.append(v)
        # 一并计算三个指标
        report = classification_report(val_targ,
                                       val_predict,
                                       target_names=remain_list,
                                       output_dict=True)

    def get(self, metrics, of_class):
        return [report[str(of_class)][metrics] for report in self.reports]

c. 新增utils.py，用于绘制训练结果图

#-*- encoding:utf-8 -*-
import pickle
import matplotlib.pyplot as plt 
import matplotlib.style as style

class MetricPlot(object):
    def __init__(self, file_pickle_path):
        self.file_pickle = file_pickle_path

    def draw_all_curves(self):
        with open(self.file_pickle + '/metric/metric.pickle', 'rb') as f:
            metric_data = pickle.load(f)
            reports = metric_data['reports']
            style.use("bmh")

            class_list = []
            for report in reports:
                class_list.extend(list(report.keys()))
            class_list = list(set(class_list))
            class_list.remove('accuracy')
            class_list.remove('macro avg')
            class_list.remove('weighted avg')

            # 每个类别分别打印precision、recall、f1
            class_num = len(class_list)
            # 每个类别分别打印precision、recall、f1
            class_num = len(class_list)
            plt.figure(figsize=(8, class_num * 4))
            for i, v in enumerate(class_list):
                plt.subplot(class_num, 1, i + 1)
                for m in {'precision', 'recall', 'f1-score'}:
                    plt.plot([report[v][m] for report in reports],
                           label='Class {0} {1}'.format(v, m))
                    plt.legend(loc='lower right')
                    plt.ylabel('Class {}'.format(v))
                    plt.title('Class {} Curves'.format(v))
            plt.show()

    def draw_metric_curves(self):
        with open(self.file_pickle + '/metric/metric.pickle', 'rb') as f:
            metric_data = pickle.load(f)
            f1 = metric_data['f1']
            recall = metric_data['recall']
            precision = metric_data['precision']

            epochs = len(f1)

            style.use("bmh")
            plt.figure(figsize=(8, 12))

            plt.subplot(3, 1, 1)
            plt.plot(range(1, epochs+1), f1, label='Val F1')
            plt.legend(loc='lower right')
            plt.ylabel('F1')
            plt.title('Validation F1 Curve')

            plt.subplot(3, 1, 2)
            plt.plot(range(1, epochs+1), recall, label='Val Recall')
            plt.legend(loc='lower right')
            plt.ylabel('Recall')
            plt.title('Validation Recall Curve')

            plt.subplot(3, 1, 3)
            plt.plot(range(1, epochs+1), precision, label='Val Precision')
            plt.legend(loc='lower right')
            plt.ylabel('Precision')
            plt.title('Validation Precision Curve')
            plt.xlabel('epoch')
            plt.show()
    def draw_history_curves(self):
        """Plot the learning curves of loss and macro f1 score 
        for the training and validation datasets.

        Args:
            history: history callback of fitting a tensorflow keras model 
        """
        with open(self.file_pickle + '/history/hist.pickle', 'rb') as f:
            history = pickle.load(f)

            loss = history['loss']
            val_loss = history['val_loss']
            accuracy = history['test_accuracy']
            val_accuracy = history['val_test_accuracy']

            epochs = len(loss)

            style.use("bmh")
            plt.figure(figsize=(8, 8)) 

            plt.subplot(2, 1, 1)
            plt.plot(range(1, epochs+1), loss, label='Training Loss')
            plt.plot(range(1, epochs+1), val_loss, label='Validation Loss')
            plt.legend(loc='upper right')
            plt.plot(range(1, epochs+1), val_loss, label='Validation Loss')
            plt.legend(loc='upper right')
            plt.ylabel('Loss')
            plt.title('Training and Validation Loss')

            plt.subplot(2, 1, 2)
            plt.plot(range(1, epochs+1), accuracy, label='Training Accuracy')
            plt.plot(range(1, epochs+1), val_accuracy, label='Validation Accuracy')
            plt.legend(loc='lower right')
            plt.ylabel('Accuracy')
            plt.title('Training and Validation Accuracy')
            plt.show()

主要的代码逻辑修改介绍完毕，下面介绍在colab上操作的流程。

三、colab操作：

1.环境部署：

# 安装tf-nightly-gpu，目前master各版本应该均可用(至少2.2.0.dev20200218亲测可用)
# 如选择发布版2.1.0，在预测阶段(重新load_model)会触发tf底层的一个bug（https://github.com/tensorflow/neural-structured-learning/issues/41），导致加载模型失败
!pip install tf-nightly-gpu
# 下载并解压tensorflow models2.1.0
!wget https://github.com/tensorflow/models/archive/v2.1.0.tar.gz
!tar xzf v2.1.0.tar.gz && rm v2.1.0.tar.gz
# 安装models需要的依赖
!pip install -r /content/models-2.1.0/official/requirements.txt

2.将修改代码上传到colab，并放置到对应的目录

# 下载数据集合和修改适配的文件
%rm -rf bert_multiclass_zh_model2-1-0/
%rm -rf process_dir
# 此处将修改好的代码传到colab上
<下载代码>
%mv /content/bert_multiclass_zh_model2-1-0/code/utils.py /content/
%mv /content/bert_multiclass_zh_model2-1-0/code/*.py /content/models-2.1.0/official/nlp/bert/
%mv /content/bert_multiclass_zh_model2-1-0/process_dir /content/

3. 为了对数据有更直观的了解，可以用下面代码对数据的分布进行观察

 1 import pandas as pd
 2 import matplotlib.pyplot as plt
 3 import matplotlib.style as style
 4 import seaborn as sns
 5 
 6 style.use("fivethirtyeight")
 7 plt.figure(figsize=(8, 12))
 8 
 9 show_data_dict = {
10     'train': pd.read_csv("/content/process_dir/train_data/train.tsv", "\t"), 
11     'dev': pd.read_csv("/content/process_dir/train_data/dev.tsv", "\t")
12 }
13 
14 # Get label frequencies in descending order
15 length = len(show_data_dict)
16 cnt = 0
17 for k, v in show_data_dict.items():
18     cnt += 1
19     plt.subplot(length, 1, cnt)
20     label_freq = v['label'].apply(lambda s: str(s)).value_counts().sort_values(ascending=False)
21     # Bar plot
22     sns.barplot(y=label_freq.index, x=label_freq.values, order=label_freq.index)
23     plt.title("{} label frequency".format(k), fontsize=14)
24     plt.xlabel("")
25     plt.xticks(fontsize=12)
26     plt.yticks(fontsize=12)
27 plt.show()

执行后可以看到数据分布情况：

 

 

 4.生成tfrecord格式的数据：

 1 !python ./models/official/nlp/bert/create_finetuning_data.py \
 2     --input_data_dir=${BASE_DIR}/train_data \
 3     --vocab_file=${BASE_DIR}/hub/vocab.txt \
 4     --train_data_output_path=${BASE_DIR}/train_data/${TASK_NAME}_train.tf_record \
 5     --eval_data_output_path=${BASE_DIR}/train_data/${TASK_NAME}_eval.tf_record \
 6     --test_data_output_path=${BASE_DIR}/train_data/${TASK_NAME}_test.tf_record \
 7     --meta_data_file_path=${BASE_DIR}/train_data/${TASK_NAME}_meta_data \
 8     --fine_tuning_task_type=classification \
 9     --max_seq_length=128 \
10     --classification_task_name=${TASK_NAME}

6. 启动训练：

 1 !python ./models/official/nlp/bert/run_classifier.py \
 2     --mode='train_and_eval' \
 3     --input_meta_data_path=${BASE_DIR}/train_data/${TASK_NAME}_meta_data \
 4     --train_data_path=${BASE_DIR}/train_data/${TASK_NAME}_train.tf_record \
 5     --eval_data_path=${BASE_DIR}/train_data/${TASK_NAME}_eval.tf_record \
 6     --test_data_path=${BASE_DIR}/train_data/${TASK_NAME}_test.tf_record \
 7     --bert_config_file=${BASE_DIR}/hub/bert_config.json \
 8     --train_batch_size=64 \
 9     --eval_batch_size=128 \
10     --test_batch_size=256 \
11     --steps_per_loop=1 \
12     --learning_rate=2e-5 \
13     --num_train_epochs=2 \
14     --model_dir=${BASE_DIR}/output \
15     --hub_module_url=https://tfhub.dev/tensorflow/bert_zh_L-12_H-768_A-12/1 \
16     --use_keras_compile_fit=True \
17     --distribution_strategy=mirrored \
18     --num_gpus=1 \
19     --save_history_path=${BASE_DIR}/history/hist.pickle \
20     --save_metric_path=${BASE_DIR}/metric/metric.pickle \
21     --model_export_path=${BASE_DIR}/export_model/save_model.h5 \
22     --test_result_dir=${BASE_DIR}/test_result/test_ret

因为在run_classifiert.py中model.fit()之前增加了一行model.summary()，所以训练开始前的日志中可见训练图层次结构等信息

整个模型参数都设置为是trainable的（默认值），笔者尝试过修改代码来freeze BERT自带的102267649个参数，但训练10个epoch后，accuracy仍达不到90%以上，感觉效果不是很好。具体原因还未深究。

7. 查看训练结果：

日志中可以查看训练中，每个batch结束后的训练loss、accuracy；每个epoch结束后的val loss、val accuracy；以及整个训练结束后的test accuracy

INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
I0223 04:38:43.759790 140623330674560 cross_device_ops.py:414] Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
I0223 04:38:43.761969 140623330674560 cross_device_ops.py:414] Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
Train for 1093 steps, validate for 313 steps
Epoch 1/2
1093/1093 [==============================] - 1161s 1s/step - loss: 0.4250 - accuracy: 0.8925 - val_loss: 0.1785 - val_accuracy: 0.9474
Epoch 2/2
1093/1093 [==============================] - 1144s 1s/step - loss: 0.1220 - accuracy: 0.9664 - val_loss: 0.1622 - val_accuracy: 0.9532
I0223 05:18:29.677855 140623330674560 run_classifier.py:420] test_accuracy: 0.951200008392334

 

查看原代码里定义的metrics（loss、accuracy），以及自定义的metrics（各类别的F1、precise、recall等）

1 import utils
2 
3 mp = utils.MetricPlot('/content/process_dir/')
4 mp.draw_history_curves()
5 mp.draw_metric_curves()
6 mp.draw_all_curves()

 

 

训练完成之后，就可以进行预测了。

首先参考create_finetune_data.py，将待预测数据转换为tfrecord格式并存储

# 预测，生成待预测数据
!python ./models-2.1.0/official/nlp/bert/do_pred_data.py \
    --input_data_dir=${PRED_DATA_DIR} \
    --vocab_file=${BASE_DIR}/model/hub_config/vocab.txt \
    --predict_data_output_path=${PRED_DATA_DIR}/${TASK_NAME}_predict.tf_record \
    --meta_data_file_path=${PRED_DATA_DIR}/${TASK_NAME}_meta_data \
    --max_seq_length=128

然后参照run_classifier.py，读取待预测tfrecord数据，加载之前训练好的模型(save_model.h5)，开始预测，将结果保存在文件中

# 开始预测
!python ./models-2.1.0/official/nlp/bert/do_predict.py \
    --input_meta_data_path=${PRED_DATA_DIR}/${TASK_NAME}_meta_data \
    --predict_data_path=${PRED_DATA_DIR}/${TASK_NAME}_predict.tf_record \
    --bert_config_file=${BASE_DIR}/model/hub_config/bert_config.json \
    --predict_batch_size=32 \
    --hub_module_url=https://tfhub.dev/tensorflow/bert_zh_L-12_H-768_A-12/1 \
    --model_export_path=${BASE_DIR}/model/saved_model/save_model.h5 \
    --predict_output_dir=${PRED_DATA_DIR}/result \
# 如果选择tf2.1.0版本，则此处会触发tf底层的一个bug（https://github.com/tensorflow/neural-structured-learning/issues/41），导致加载模型失败

最终训练出来的文件内容是这样的：第一列是标签的编号，第二列是标签，由于是串行训练，因此输出的标签和输入的预测文本是一一对齐的。

1    poi
4    webfic
7    science
0    human
0    human
0    human
1    poi
7    science

至此，整个finetuneing训练+预测过程全部结束。欢迎拍砖指导~~~~

稍后再研究一下如何部署到生产环境……