使用 ERNIE 完成标注任务
semal_smile 发布于2019-11 浏览:3071 回复:0
2
收藏
最后编辑于2020-10

## 摘要
ERNIE 是 BERT 在 paddlepaddle 框架下的实现和优化。由于是百度的平台,对于中文有更好的性能。基于 paddlepaddle 框架训练模型和部署预测,速度上较 Tensorflow 也快了许多。本文介绍如何使用 ERNIE 训练标注任务的模型并部署相关服务。

## 背景
基于 Tensorflow 的 BERT 实现在使用各种加速方法下,部署服务的速度依然不太能满足生产环境的需求。在要求更快的预测速度的需求下,paddlepaddle 框架下的 ERNIE 已经在百度的生产环境下使用,并有着不错的迭代速度。

### paddlepaddle
paddlepaddle 的编程称为 Fluid 编程。Fluid 使用 Tensor 定义数据,Operator 定义对数据的操作,Operator 进一步封装在 fluid.layer 和 fluid.net 中。计算图的结构存储在 Program 中,通过 Executor 执行 Program 训练模型。模型的保存分为保存参数和保存预测模型两种,save_params 和 save_persistables 只保存参数,save_inference_model 保存参数和计算图,并指定输入输出节点用于预测。load_params、load_persistables 和 load_inference_model 用于加载训练好的模型。

## 方法
本文使用的是 ERNIE_stable-1.0.1 版本,在 python 2.7.16 下进行的测试,paddlepaddle 的版本为 1.5.1。

### 标注数据准备
ERNIE 的标注数据文件夹下包含四个文件:train.tsv、dev.tsv、test.tsv 和 label_map.json。*.tsv 文件内容如下所示,文本列和标记列的元素之间用“\002”分隔。
```
text_a label
阜新市细河区五福园饭店 OOOOOOB-REPRESENTI-REPRESENTI-REPRESENTOO
苏州天炫动画设计有限公司 OOB-REPRESENTI-REPRESENTI-REPRESENTI-REPRESENTOOOOOO
凤冈县壹等壹不锈钢经销部 OOOB-REPRESENTI-REPRESENTI-REPRESENTOOOOOO
```
label_map.json 定义了标注集对应的数字编号。
```json
{
"O": 0,
"B-REPRESENT": 1,
"I-REPRESENT": 2
}
```

### 保存预测模型
ERNIE 提供的训练脚本 run_sequence_labeling.py 是通过 save_persistable 保存的模型,只保存了模型参数,需要额外写脚本重新定义计算图结构,然后用保存的模型参数对计算图赋值,然后调用 save_inference_model 保存成预测需要的模型。
```python
def convert_persistent2infer_model():
model_path = 'data/pretrained/ERNIE_stable-1.0.1'
ernie_config_path = '%s/ernie_config.json' % model_path
ernie_config = ErnieConfig(ernie_config_path)
ernie_config.print_config()
use_cuda = False
max_seq_len = 50
save_inference_model_path = 'trade-name-infer-model'
init_checkpoint = 'checkpoints-trade-name/step_6000'

args.max_seq_len = max_seq_len
args.use_fp16 = False
args.num_labels = 3
# 构建计算图结构
predict_prog = fluid.Program()
predict_startup = fluid.Program()
with fluid.program_guard(predict_prog, predict_prog):
with fluid.unique_name.guard():
predict_pyreader, infer_labels, feed_target_names = create_model(
args,
pyreader_name='predict_reader',
ernie_config=ernie_config,
is_prediction=True)
predict_prog = predict_prog.clone(for_test=True)
place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()
exe = fluid.Executor(place)
exe.run(predict_startup)
# 将训练好的参数赋值给计算图
fluid.io.load_persistables(exe, init_checkpoint, predict_prog)
_, ckpt_dir = os.path.split(init_checkpoint.rstrip('/'))
dir_name = ckpt_dir + '_inference_model'
model_path = os.path.join(save_inference_model_path, dir_name)
print("save inference model to %s" % model_path)
# 保存预测模型,feed_target_names 指定输入节点,[infer_labels] 指定输出节点
fluid.io.save_inference_model(
model_path,
feed_target_names,
[infer_labels],
exe,
main_program=predict_prog)
```

### 使用预测模型进行预测
利用 load_inference_model 导入预测模型,可以获得计算图 Program 和输入输出节点。通过喂给输入节点数据,使用 Executor 执行 Program 就能从输出节点得到预测结果。这里需要注意的是输入文本需要使用 [CLS] 和 [SEP] 标记包住,批量预测的时候要将数据用 0 扩展到等长。
```python
def predict():
model_path = '/qichacha2/gongzh/ERNIE/trade-name-infer-model/step_2000_inference_model'
# model_path = '/qichacha2/gongzh/ERNIE/trade-name-test-save-infer-model'
print("load inference model from %s" % model_path)
# place = fluid.CPUPlace()
place = fluid.CUDAPlace(0)
exe = fluid.Executor(place)
# 导入预测模型
infer_program, feed_target_names, fetch_targets = fluid.io.load_inference_model(model_path, exe)
# print(probs)
# print(feed_target_names)

max_seq_len = 26

src_ids = feed_target_names[0]
sent_ids = feed_target_names[1]
pos_ids = feed_target_names[2]
input_mask = feed_target_names[3]

tokenizer = FullTokenizer('data/pretrained/ERNIE_stable-1.0.1/vocab.txt')
text1 = '上海阅文信息技术有限公司'
text2 = '思南县亿家诚汽贸有限公司'
text3 = '北安市禾思辉聚电子商务有限公司'
text4 = '遵义智讯商贸有限公司'
text5 = '贵州珏摩网络科技有限公司'
texts = [text1, text2, text3, text4, text5]
tokens_list = []
for text in texts:
tokens = tokenizer.tokenize(text)
tokens = ["[CLS]"] + tokens + ["[SEP]"]
tokens_list.append(tokens)
max_seq_len = max(max_seq_len, len(tokens))
token_ids, position_ids, _input_mask = [], [], []
for tokens in tokens_list:
_token_ids = tokenizer.convert_tokens_to_ids(tokens)
_position_ids = list(range(len(_token_ids)))
__input_mask = [1.] * len(_token_ids)
_token_ids = _token_ids[:max_seq_len] + [0] * (max_seq_len - len(_token_ids))
_position_ids = _position_ids[:max_seq_len] + [0] * (max_seq_len - len(_position_ids))
__input_mask = __input_mask[:max_seq_len] + [0.] * (max_seq_len - len(__input_mask))
token_ids.extend(_token_ids)
position_ids.extend(_position_ids)
_input_mask.extend(__input_mask)

batch_size = len(texts)
src_ids_data = np.array(token_ids).reshape((batch_size, max_seq_len, 1))
sent_ids_data = np.zeros((batch_size, max_seq_len, 1), dtype=np.int64)
pos_ids_data = np.array(position_ids).reshape((batch_size, max_seq_len, 1))
input_mask_data = np.array(_input_mask, dtype=np.float32).reshape((batch_size, max_seq_len, 1))
start = time.clock()
output = exe.run(
infer_program,
feed={
src_ids: src_ids_data,
sent_ids: sent_ids_data,
pos_ids: pos_ids_data,
input_mask: input_mask_data
},
fetch_list=fetch_targets)
label_ids_list = output[0].reshape((batch_size, max_seq_len,))
for i, label_ids in enumerate(label_ids_list):
print(label_ids)
tokens = tokens_list[i][1:-1]
trade_name = ''
for j, label_id in enumerate(label_ids[1:len(tokens) - 1]):
if label_id > 0:
trade_name += tokens[j]
print(trade_name)
consume = time.clock() - start
print(consume)
```

## 结果
GPU 下单次预测的延迟在 30ms 到 50ms 之间,3G 显存下可以批量 100 左右的数据同时预测而不影响延迟。ERNIE 对从公司名中提取商号的标注任务性能可以达到 96% 以上的正确率和召回率,较 BERT 还要高出两个百分点。

## 讨论
ERNIE 不论在延迟还是性能上都较 BERT 有更好的结果,有很好的使用前景,建议在未来的相关自然语言处理任务中可以以 ERNIE 作为主要模型训练和预测引擎。

收藏
点赞
2
个赞
快速回复
TOP
切换版块