tensorflow(7)利用tensorflow/serving实现BERT模型部署

  本文将会详细介绍如何使用tensorflow/serving来实现BERT模型的部署及预测。
  我们以Github上的bertNER为例，该项目使用BERT+Bi-LSTM+CRF结构实现中文序列标注，对BERT进行微调，并且提供了模型训练、模型预测的办法。本文将在此基础上，将模型训练后生成的ckpt文件转化为pb文件（tensorflow/serving支持的部署文件格式），并使用tensorflow/serving来实现BERT模型的部署及预测。

如何将ckpt文件转化为pb文件？

  该项目使用BERT+Bi-LSTM+CRF结构实现中文序列标注，对BERT进行微调,因此模型无疑是复杂的。
  我们先使用上述项目对人民日报命名实体识别数据集进行训练，得到ckpt(训练10个epoch，在test数据集上的F1值为92.55%)文件如下:

$ tree ckpt/
ckpt/
├── checkpoint
├── ner.ckpt-1141.data-00000-of-00001
├── ner.ckpt-1141.index
├── ner.ckpt-1141.meta
├── ner.ckpt-1467.data-00000-of-00001
├── ner.ckpt-1467.index
├── ner.ckpt-1467.meta
├── ner.ckpt-1630.data-00000-of-00001
├── ner.ckpt-1630.index
├── ner.ckpt-1630.meta
├── ner.ckpt-815.data-00000-of-00001
├── ner.ckpt-815.index
├── ner.ckpt-815.meta
├── ner.ckpt-978.data-00000-of-00001
├── ner.ckpt-978.index
└── ner.ckpt-978.meta

0 directories, 16 files

  首先我们需要观察模型的结构。在文章tensorflow(4)使用tensorboard查看ckpt和pb图结构中，我们知道了如何使用tensorboard查看ckpt图结构。python代码如下：

import tensorflow as tf
from tensorflow.summary import FileWriter

sess = tf.Session()
tf.train.import_meta_graph("ner.ckpt-1630.meta")
FileWriter("log", sess.graph)
sess.close()

在ckpt目录下运行该脚本，生成log文件夹，使用命令：tensorboard --logdir=log，在浏览器中localhost:6006，即可查看模型图结构，如下图：

从上面的模型结构图中，我们可以知道，该模型的输入为input_ids、input_mask、segment_ids、dropout，再通过代码可知输出为logits.在知道了模型输入、输出后，利用下面的脚本可以方便地将ckpt文件转换为pb文件，代码如下：

import os
import pickle
import tensorflow as tf
from utils import create_model, get_logger
from train import FLAGS, load_config
from tensorflow.python import saved_model
from model import Model

config = load_config(FLAGS.config_file)
logger = get_logger(FLAGS.log_file)
# limit GPU memory
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
with open(FLAGS.map_file, "rb") as f:
    tag_to_id, id_to_tag = pickle.load(f)

export_path = "~/bertNER/deploy_models/people_daily/1"
with tf.Session(config=tf_config) as sess:
    model = create_model(sess, Model, FLAGS.ckpt_path, config, logger)
    saved_model.simple_save(session=sess,
                            export_dir=export_path,
                            inputs={"input_ids": model.input_ids, "input_mask": model.input_mask,
                                    "segment_ids": model.segment_ids, "dropout": model.dropout},
                            outputs={"logits": model.logits})

查看deploy_models，结构如下：

$ tree deploy_models/
deploy_models/
└── people_daily
    └── 1
        ├── saved_model.pb
        └── variables
            ├── variables.data-00000-of-00001
            └── variables.index

3 directories, 3 files

利用tensorflow/serving部署模型

  有了pb文件，我们借助文章tensorflow(6)利用tensorflow/serving实现模型部署及预测和tensorflow(5)将ckpt转化为pb文件并利用tensorflow/serving实现模型部署及预测的帮助，利用tensorflow/serving的Docker镜像进行模型部署，命令如下：

docker run -t --rm -p 8531:8501 -v "$path/bertNER/deploy_models/people_daily:/models/people_daily" -e MODEL_NAME=people_daily tensorflow/serving:1.14.0

$path为bertNER所在目录。查看模型是否部署成功：

curl http://192.168.1.193:8531/v1/models/people_daily

输出结果为：

{
    "model_version_status": [
        {
            "version": "1",
            "state": "AVAILABLE",
            "status": {
                "error_code": "OK",
                "error_message": ""
            }
        }
    ]
}

模型已经成功部署。

如何进行模型预测？

  在之前的系列文章中，已经介绍了使用调用tensorflow/serving的服务。在这里，我们只需要输入input_ids、input_mask、segment_ids，而这个笔者已经在项目中给出现成的函数，我们直接使用就可以了，以下是模型预测的一个示例脚本：

# -*- coding: utf-8 -*-
import pickle
import requests
import numpy as np
import tensorflow as tf

from train import FLAGS
from loader import input_from_line


# limit GPU memory
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True


with open("maps.pkl", "rb") as f:
    tag_to_id, id_to_tag = pickle.load(f)

sentence = "日本漫画家井上雄彦宣布《灌篮高手》将拍成电影。"

string, segment_ids, ids, mask, label_ids = input_from_line(sentence, FLAGS.max_seq_len, tag_to_id)

segment_ids = segment_ids.tolist()[0]
ids = ids.tolist()[0]
mask = mask.tolist()[0]

tensor = {"instances": [{"segment_ids": segment_ids, "input_ids": ids, "input_mask": mask, "dropout": 1}]}

req = requests.post("http://192.168.1.193:8531/v1/models/people_daily:predict", json=tensor)

res = {}
if req.status_code == 200:
    t = np.asarray(req.json()['predictions'][0]).argmax(axis=1)
    tags = [id_to_tag[_] for _ in t]
    print(tags)

输出的预测序列为：

'[CLS]', 'B-LOC', 'I-LOC', 'O', 'O', 'O', 'B-PER', 'I-PER', 'I-PER', 'I-PER', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', '[SEP]'

对比预测序列和文字，我们不难发现，日本是LOC，井上雄彦是PER。

总结

  本文虽然将BERT模型如何在tensorflow/serving讲得比较清晰，但实际笔者的探索过程却是非常曲折的。
  当然，文章tensorflow(6)利用tensorflow/serving实现模型部署及预测和tensorflow(5)将ckpt转化为pb文件并利用tensorflow/serving实现模型部署及预测可以帮助读者更好地理解本文的内容。
  感谢大家的阅读~