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

水深无声 2023-01-03 11:25 375阅读 0赞

本文将会详细介绍如何使用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图结构][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，即可查看模型图结构，如下图：  
![模型结构图][20210107214440783.png_pic_center]

从上面的模型结构图中，我们可以知道，该模型的输入为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_6_tensorflow_serving]和[tensorflow(5)将ckpt转化为pb文件并利用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_6_tensorflow_serving]和[tensorflow(5)将ckpt转化为pb文件并利用tensorflow/serving实现模型部署及预测][tensorflow_5_ckpt_pb_tensorflow_serving]可以帮助读者更好地理解本文的内容。  
  感谢大家的阅读~

[tensorflow_4_tensorboard_ckpt_pb]: https://blog.csdn.net/jclian91/article/details/109508275
[20210107214440783.png_pic_center]: /images/20221119/205ec609fe504a149c3f49adb718afc2.png
[tensorflow_6_tensorflow_serving]: https://blog.csdn.net/jclian91/article/details/112298586
[tensorflow_5_ckpt_pb_tensorflow_serving]: https://blog.csdn.net/jclian91/article/details/109521659