Google Bert 框架训练、验证、推断和导出简单说明
实习的这段期间，在公司做了不少 NLP 分类任务，歧义车系判断、字词重复纠错等等，期间有用过 Google 开源的 Bert 框架，也用过公司大佬制作的 T5 模型。但无论使用什么，起手 Bert 仿佛已经成为了一种“本能”（笑哭.jpg），Bert NLP 算法工程师的至交好友。
写这篇博客的目的一是为了记录先前工作的经验，此外也简单地介绍一下如何使用 Google 官方开源的 Bert 框架，因为目前很少有博客会讲如何将 Bert 训练得到的 checkpoint 转换为 savedModel，这何尝不是一种遗憾，因此我打算将这遗憾填补。对于第一次使用 Bert 框架的读者，我建议从头开始看，若已经有熟练的使用经验，只想了解如何导出模型和使用，可以直接跳到模型导出和使用。
【项目地址】：https://github.com/clvsit/bert-simple-use
准备工作
首先，到 GitHub 上 clone Bert 源代码。
- google-research/bert：https://github.com/google-research/bert
然后下载预训练模型，本博客的案例用的是中文数据，因此下载 BERT-Base, Chinese（如果下载很慢，可以尝试到网盘或其他资源站点去下载）。
【数据】：案例使用的数据选择外卖情感极性评价数据集，这个数据集是我从 CSDN 的下载中找到的。
- 积极：
很快，好吃，味道足，量大
没有送水没有送水没有送水
非常快，态度好。
方便，快捷，味道可口，快递给力
菜味道很棒！送餐很及时！
今天师傅是不是手抖了，微辣格外辣！
送餐快,态度也特别好,辛苦啦谢谢
超级快就送到了，这么冷的天气骑士们辛苦了。谢谢你们。麻辣香锅依然很好吃。
经过上次晚了2小时，这次超级快，20分钟就送到了……
最后五分钟订的，卖家特别好接单了，谢谢。
- 消极：
菜品质量好，味道好，就是百度的问题，总是用运力原因来解释，我也不懂这是什么原因，晚了三个小时呵呵厉害吧！反正订了就退不了，只能干等……
分量还可以……就是有点没特色……下回不吃啦
没什么味道，送来的晚凉了
送餐送错，还狡辩不给补偿，送餐时间3个小时，百度送餐员更是素质卑劣，额外还要加收17元的外送费。百度客服也没用，也没有解决！
完全不值得信任，出尔反而
最差餐厅，没有之一
感觉不太好吃，价格贵。但是百度外卖的送餐人员态度很好！
没发票，乱收费，订单没到就被完成了
量很大，但是味道真的一般。等了一个多小时才送到，达到了超时赔付的时间，感觉百度超时赔付就是个摆设
头一天看同事点的三份菜就一大盆，今天点了四份菜才一小碗……差太大吧，辣椒花椒太多，非常油腻……
然后，我们统计所有数据的句长，其中最长的句子长度为 463。
可以看到句子长度的分布很不均匀，这里我们可以将 max_seq_length 设置成 465（加上 [CLS] 和 [SEQ]），对于长度不足 465 的句子通过 padding 填充，缺点在于占用更多的空间和花费更多的训练时间。实际上句长较长的那部分数据属于长尾数据，我们可以将其抛弃。最终得到的是句子长度小于等于 126 的数据，max_seq_length 设置为 128。
接着将上述两部分的数据处理成 BERT 的输入数据形式，在处理成 csv 数据格式时需要注意“,”（英文逗号）是否在原文中出现，可以看到这份数据不够“干净”，里面存在大量的错别字以及错用的标点符号，因此有两种处理方式：
- 将英文逗号替换为中文逗号，然后用英文逗号作为输入文本和标签的分隔符。
- 使用文章中没有出现的符号作为分隔符。
这里采用第二种方法，使用“&”作为文本和标签的分隔符。
吃的挺好的,以后还会点别的&1
送餐比之前快了不少呀，是只有我这么想么。味道很赞。&1
分量够，味道可以，送一次性手套和餐巾纸&1
不错，常客了，肘子的好吃,送货也快&1
煎饼很好吃！送餐很快！&1
皮蛋粥快咸死了，不好吃！不过包装值得表扬！&0
还行吧，因为送来时有点凉了，等的有点久&0
味儿还行，就是油太大了！&0
加了一份米饭。打开一看。两份顶一份。太坑了。2份饭不够吃。&0
沙拉恶心死了都成泥了快递竟然还没有零钱&0
2:30送到,小伙伴们,看着办吧&0
处理完数据之后，我们再将 bert 的代码和预训练模型 chinese_L-12_H-768_A-12 放到同一个项目下，整个项目结构如下所示。
【项目结构】：
model/
vocab.txt
bert_model.ckpt.meta
bert_model.ckpt.index
bert_model.ckpt.data-00000-of-00001
bert_config.json
data/
train.csv
dev.csv
test.csv
bert/
...
output/
export/
train.sh
predict.sh
predict.py
export.sh
训练和验证
使用 BERT 框架进行模型训练非常简单，我们要做的就是修改（1）数据读取部分（2）模型配置部分。
（1）数据读取部分：找到 run_classifier.py 文件中的 DataProcessor(object) 类，我们要做的就是继承这个类，用来处理我们自己的数据。
我们可以直接复制已有的 DataProcessor，然后修改下相应的代码。
- 修改读取文件的路径和名称。
def get_train_examples(self, data_dir):
return self._create_examples(
self._read_tsv(os.path.join(data_dir, "train.csv")), "train")
def get_dev_examples(self, data_dir):
return self._create_examples(
self._read_tsv(os.path.join(data_dir, "dev.csv")), "dev")
def get_test_examples(self, data_dir):
return self._create_examples(
self._read_tsv(os.path.join(data_dir, "test.csv")), "test")
将读取 train、dev 和 test 函数中的文件名称修改成我们的文件名称，其余都可以不用改动。
- 修改标签函数。
def get_labels(self):
return ["0", "1"]
因为是一个二分类任务（积极和消极），因此可以让 get_labels() 函数直接返回 0 和 1。
- 修改文件读取函数。
def _read_tsv(cls, input_file, quotechar=None):
with tf.gfile.Open(input_file, "r") as f:
reader = csv.reader(f, delimiter="&", quotechar=quotechar)
lines = []
for line in reader:
lines.append(line)
return lines
因为我们的数据是以“&”作为分隔符，因此在这需要将 delimiter="\t" 修改为 delimiter="&"。
- 修改 _create_examples() 函数。
def _create_examples(self, lines, set_type):
examples = []
for (i, line) in enumerate(lines):
guid = "%s-%s" % (set_type, i)
if set_type == "test":
text_a = tokenization.convert_to_unicode(line[0])
label = "0"
else:
text_a = tokenization.convert_to_unicode(line[0])
label = tokenization.convert_to_unicode(line[1])
if label not in ["0", "1"]:
continue
examples.append(
InputExample(guid=guid, text_a=text_a, text_b=None, label=label))
return examples
_create_examples() 函数是修改的重点，在这里处理数据中各字段的读取，因为这是一个单句子任务，因此我们只需要 text_a 即可。
【完整代码】：
class EmotionProcessor(DataProcessor):
def get_train_examples(self, data_dir):
return self._create_examples(
self._read_tsv(os.path.join(data_dir, "train.csv")), "train")
def get_dev_examples(self, data_dir):
return self._create_examples(
self._read_tsv(os.path.join(data_dir, "dev.csv")), "dev")
def get_test_examples(self, data_dir):
return self._create_examples(
self._read_tsv(os.path.join(data_dir, "test.csv")), "test")
def get_labels(self):
return ["0", "1"]
def _read_tsv(cls, input_file, quotechar=None):
with tf.gfile.Open(input_file, "r") as f:
reader = csv.reader(f, delimiter="&", quotechar=quotechar)
lines = []
for line in reader:
lines.append(line)
return lines
def _create_examples(self, lines, set_type):
examples = []
for (i, line) in enumerate(lines):
guid = "%s-%s" % (set_type, i)
if set_type == "test":
text_a = tokenization.convert_to_unicode(line[0])
label = "0"
else:
text_a = tokenization.convert_to_unicode(line[0])
label = tokenization.convert_to_unicode(line[1])
if label not in ["0", "1"]:
continue
examples.append(
InputExample(guid=guid, text_a=text_a, text_b=None, label=label))
return examples
最后，将新创建的 EmotionProcessor 加入到 main() 函数的 processors 中。
processors = {
"emotion": EmotionProcessor,
}
【注意】：这里的 emotion 需要小写，因为 BERT 在读取 task_name 时进行了小写处理。
task_name = FLAGS.task_name.lower()
if task_name not in processors:
raise ValueError("Task not found: %s" % (task_name))
processor = processors[task_name]()
（2）模型配置部分：run_classifier.py 文件是一个命令行调用脚本文件，如果是在 linux 系统上，我们可以编写 bash 脚本，在这将训练和验证一起完成。
#! /bin/bash
export MODEL_DIR=model
export DATA_DIR=data
CUDA_VISIBLE_DEVICES=1 python bert/run_classifier.py \
--task_name=Emotion \
--do_train=true \
--do_eval=true \
--data_dir=$DATA_DIR \
--vocab_file=$MODEL_DIR/vocab.txt \
--bert_config_file=$MODEL_DIR/bert_config.json \
--init_checkpoint=$MODEL_DIR/bert_model.ckpt \
--max_seq_length=128 \
--train_batch_size=32 \
--learning_rate=2e-5 \
--num_train_epochs=2.0 \
--output_dir=output
【注意】：
- 请根据自己机子的配置设置
train_batch_size大小，以及 bert_config.json 中的配置内容。 - CUDA_VISIBVLE_DEVICES 指定要使用的显卡，如果只有一张显卡，设置为 0，即
CUDA_VISIBLE_DIVICES=0。
如果没有问题的话，我们就以 bert 默认的配置进行训练和验证。在控制台输入：
sh train.sh
等待一段时间后，直到训练完成后，我们可以在控制台中看到模型在验证集上的准确率。
模型推断
在完成模型训练后，我们可以在 output 目录下看到模型训练和验证阶段的记录和结果。
eval/
train.tf_record
model.ckpt-624.meta
model.ckpt-624.index
model.ckpt-624.data-00000-of-00001
model.ckpt-0.meta
model.ckpt-0.index
model.ckpt-0.data-00000-of-00001
graph.pbtxt
events.out.tfevents.xxx
eval_results.txt
eval.tf_record
checkpoint
其中 train.tf_record 和 eval.tf_record 是我们训练和验证数据集，eval_results.txt 是验证的结果，而 model.ckpt-624 是训练完成的模型文件，也是我们推断时指定的模型。
在项目目录下创建 predict.sh 文件。
#! /bin/bash
export MODEL_DIR=model
export DATA_DIR=data
CUDA_VISIBLE_DEVICES=1 python bert/run_classifier.py \
--task_name=Emotion \
--do_predict=true \
--data_dir=$DATA_DIR \
--vocab_file=$MODEL_DIR/vocab.txt \
--bert_config_file=$MODEL_DIR/bert_config.json \
--init_checkpoint=output/model.ckpt-624 \
--output_dir=output
相比 train.sh 文件，将 do_train 和 do_eval 修改为 do_predict，并指定 init_checkpoint 为我们训练完成的模型文件。如果你想将推断结果存放到其他目录下，可以修改 output_dir 参数值。在这，我仍然将推断的结果存储到 output 目录下，此时可以看到多了 test_results.tsv 文件。
0.011303517 0.98869646
0.9418804 0.058119625
0.04432816 0.9556718
0.011732221 0.98826784
0.029930793 0.9700693
0.012002373 0.98799765
0.09023312 0.90976685
0.017287388 0.9827126
0.03125599 0.968744
0.015658164 0.9843418
该文件记录模型对各标签的预测概率值，例如第一条消极的概率为 0.011303517，积极的概率为 0.98869646，模型认为第一条评论是积极的。查看 test.csv 文件的第一条评论“菜量很大，味道也不错，师傅速度很快，好评～”，模型的预测是正确的。
模型导出和使用
虽然我们可以直接使用 predict.sh 去做预测，但问题是我们需要将预测的数据做成 csv 文件，然后启动 predict.sh，能不能做成函数的形式，把输入传给一个函数，然后得到相应的结果？当然可以！
首先，我们需要在 run_classifier.py 文件中新增导出的代码。
- 在代码的 flags 区域加上
export_dir和do_export。
flags.DEFINE_string(
"export_dir", None,
"The dir where the exported model will be written.")
flags.DEFINE_bool(
"do_export", False, "Whether to export the model.")
- 然后创建
serving_input_fn()函数。
def serving_input_fn():
input_ids = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='input_ids')
input_mask = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='input_mask')
segment_ids = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='segment_ids')
input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn({
'label_ids': tf.constant(0, tf.int32),
'input_ids': input_ids,
'input_mask': input_mask,
'segment_ids': segment_ids
})()
return input_fn
- 接着在 main() 函数中修改如下几处代码。
if not FLAGS.do_train and not FLAGS.do_eval and not FLAGS.do_predict and not FLAGS.do_export:
raise ValueError(
"At least one of `do_train`, `do_eval` or `do_predict' must be True.")
- 最后，在 main() 函数的尾部加上导出相关的代码。
if FLAGS.do_export:
estimator._export_to_tpu = False
estimator.export_savedmodel(FLAGS.export_dir, serving_input_fn)
完成代码的修改后，接着把训练好的模型转换成 savedModel 形式，开始编写 export.sh 文件。
#! /bin/bash
export MODEL_DIR=model
export DATA_DIR=data
CUDA_VISIBLE_DEVICES=1 python bert/run_classifier.py \
--task_name=Emotion \
--do_export=true \
--data_dir=$DATA_DIR \
--vocab_file=$MODEL_DIR/vocab.txt \
--bert_config_file=$MODEL_DIR/bert_config.json \
--init_checkpoint=output/model.ckpt-624 \
--output_dir=output
--export_dir=export
运行 sh export.sh 命令，等待片刻后，看到控制台输出：
SavedModel written to: exported/temp-b'1590300832'/saved_model.pb
export 目录下多了 temp-b’1590300779’ 和 1590300832 的两个目录，其中 1590300832 是导出的 saveModel 以时间戳命名。
最后，让我们编写一个简单的 predict.py 脚本文件。
import tensorflow as tf
from bert import tokenization
def convert_single_example(query, label_list, max_seq_length, tokenizer):
"""Converts a single `InputExample` into a single `InputFeatures`."""
label_map = {
}
for (i, label) in enumerate(label_list):
label_map[label] = i
tokens_a = tokenizer.tokenize(query)
if len(tokens_a) > max_seq_length - 2:
tokens_a = tokens_a[0:(max_seq_length - 2)]
tokens = []
segment_ids = []
tokens.append("[CLS]")
segment_ids.append(0)
for token in tokens_a:
tokens.append(token)
segment_ids.append(0)
tokens.append("[SEP]")
segment_ids.append(0)
input_ids = tokenizer.convert_tokens_to_ids(tokens)
input_mask = [1] * len(input_ids)
while len(input_ids) < max_seq_length:
input_ids.append(0)
input_mask.append(0)
segment_ids.append(0)
return {
"input_ids": input_ids,
"segment_ids": segment_ids,
"input_mask": input_mask
}
if __name__ == '__main__':
label_list = [0, 1]
predict_fn = tf.contrib.predictor.from_saved_model("exported/1590300832/")
tokenizer = tokenization.FullTokenizer(vocab_file="model/vocab.txt", do_lower_case=True)
feature = convert_single_example("菜量很大，味道也不错，师傅速度很快，好评～", label_list, 128, tokenizer)
prediction = predict_fn({
"input_ids": [feature['input_ids']],
"segment_ids": [feature['segment_ids']],
"input_mask": [feature['input_mask']]
})
probabilities = prediction["probabilities"]
label = label_list[probabilities.argmax()]
print(probabilities)
print(label)
其中，convert_single_example() 函数可以直接从 run_classifier.py 文件中拷贝。我们要做的实际上就是读取 saveModel 文件以及整理输入数据格式。
- 读取 saveModel：注意替换成自己模型的名称哦。
predict_fn = tf.contrib.predictor.from_saved_model("exported/1590300832/")
- 整理输入数据格式。
tokenizer = tokenization.FullTokenizer(vocab_file="model/vocab.txt", do_lower_case=True)
feature = convert_single_example("菜量很大，味道也不错，师傅速度很快，好评～", label_list, 128, tokenizer)
- 将输入数据传给预测函数，得到预测结果。
prediction = predict_fn({
"input_ids": [feature['input_ids']],
"segment_ids": [feature['segment_ids']],
"input_mask": [feature['input_mask']]
})
至此，关于 Bert 框架训练、验证、推断和导出的简单说明告一段落，如有错误请各位读者指出，不胜感激。