NLP之BERT分类模型部署提供服务_bert模型 部署

在我们使用bert预分类模型微调之后（可以参考我前面写的文章），需要对项目进行支持，那就需要分类模型落地提供服务，这篇文章介绍python调用bert模型，提供服务。
参考：https://github.com/xmxoxo/BERT-train2deploy
在后期部署的时候，需要一个label2id的字典，所以要在训练的时候就保存起来，在

convert_single_example这个方法里增加一段：
  #--- save label2id.pkl ---
  output_label2id_file = os.path.join(FLAGS.output_dir, "label2id.pkl")
  if not os.path.exists(output_label2id_file):
    with open(output_label2id_file,'wb') as w:
      pickle.dump(label_map,w)
1
2
3
4
5
6

这样训练后就会生成这个文件了。

1、转换模型

在训练bert模型之后会得到一个output文件夹，里面是tf的checkout文件，模型是.ckpt的文件格式,文件比较大，并且有多个文件：

可以看到，模板文件非常大，大约有好几个G（个数和大小由训练的语料和参数决定）。 后面使用的模型服务端，使用的是.pb格式的模型文件，所以需要把生成的ckpt格式模型文件转换成.pb格式的模型文件。提供了一个转换工具:freeze_graph.py，使用如下：

usage: freeze_graph.py [-h] -bert_model_dir BERT_MODEL_DIR -model_dir
                       MODEL_DIR [-model_pb_dir MODEL_PB_DIR]
                       [-max_seq_len MAX_SEQ_LEN] [-num_labels NUM_LABELS]
                       [-verbose]
1
2
3
4

这里要注意的参数是：

• model_dir 就是训练好的.ckpt文件所在的目录
• max_seq_len 要与原来一致；
• num_labels 是分类标签的个数,本例中是7个

python freeze_graph.py \
    -bert_model_dir $BERT_BASE_DIR \
    -model_dir $TRAINED_CLASSIFIER/$EXP_NAME \
    -max_seq_len 128 \
    -num_labels 7
1
2
3
4
5

执行成功后可以看到在model_dir目录会生成一个classification_model.pb 文件。 转为.pb格式的模型文件，同时也可以缩小模型文件的大小,可以看到转化后的模型文件大约是390M。

-rw-r--r-- 1 root root  409344978 Jun  3 11:01 classification_model.pb
1

附 freeze_graph.py：

# -*- coding: utf-8 -*-
"""
Created on Sun Apr 28 10:20:04 2019

@author: chenyang
"""

#import contextlib
import json
import os
from enum import Enum
from termcolor import colored
import sys
import modeling
import logging
import tensorflow as tf
import argparse


def set_logger(context, verbose=False):
    if os.name == 'nt':  # for Windows
        return NTLogger(context, verbose)

    logger = logging.getLogger(context)
    logger.setLevel(logging.DEBUG if verbose else logging.INFO)
    formatter = logging.Formatter(
        '%(levelname)-.1s:' + context + ':[%(filename).3s:%(funcName).3s:%(lineno)3d]:%(message)s', datefmt=
        '%m-%d %H:%M:%S')
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.DEBUG if verbose else logging.INFO)
    console_handler.setFormatter(formatter)
    logger.handlers = []
    logger.addHandler(console_handler)
    return logger


class NTLogger:
    def __init__(self, context, verbose):
        self.context = context
        self.verbose = verbose

    def info(self, msg, **kwargs):
        print('I:%s:%s' % (self.context, msg), flush=True)

    def debug(self, msg, **kwargs):
        if self.verbose:
            print('D:%s:%s' % (self.context, msg), flush=True)

    def error(self, msg, **kwargs):
        print('E:%s:%s' % (self.context, msg), flush=True)

    def warning(self, msg, **kwargs):
        print('W:%s:%s' % (self.context, msg), flush=True)

def create_classification_model(bert_config, is_training, input_ids, input_mask, segment_ids, labels, num_labels):

    #import tensorflow as tf
    #import modeling

    # 通过传入的训练数据，进行representation
    model = modeling.BertModel(
        config=bert_config,
        is_training=is_training,
        input_ids=input_ids,
        input_mask=input_mask,
        token_type_ids=segment_ids,
    )

    embedding_layer = model.get_sequence_output()
    output_layer = model.get_pooled_output()
    hidden_size = output_layer.shape[-1].value

    output_weights = tf.get_variable(
        "output_weights", [num_labels, hidden_size],
        initializer=tf.truncated_normal_initializer(stddev=0.02))

    output_bias = tf.get_variable(
        "output_bias", [num_labels], initializer=tf.zeros_initializer())

    with tf.variable_scope("loss"):
        if is_training:
            # I.e., 0.1 dropout
            output_layer = tf.nn.dropout(output_layer, keep_prob=0.9)

        logits = tf.matmul(output_layer, output_weights, transpose_b=True)
        logits = tf.nn.bias_add(logits, output_bias)
        probabilities = tf.nn.softmax(logits, axis=-1)
        log_probs = tf.nn.log_softmax(logits, axis=-1)

        if labels is not None:
            one_hot_labels = tf.one_hot(labels, depth=num_labels, dtype=tf.float32)

            per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1)
            loss = tf.reduce_mean(per_example_loss)
        else:
            loss, per_example_loss = None, None
    return (loss, per_example_loss, logits, probabilities)


def init_predict_var(path):
    label2id_file = os.path.join(path, 'label2id.pkl')
    if os.path.exists(label2id_file):
        with open(label2id_file, 'rb') as rf:
            label2id = pickle.load(rf)
            id2label = {value: key for key, value in label2id.items()}
            num_labels = len(label2id.items())
    return num_labels, label2id, id2label



def optimize_class_model(args, logger=None):
    if not logger:
        logger = set_logger(colored('CLASSIFICATION_MODEL, Lodding...', 'cyan'), args.verbose)
        pass
    try:
        # 如果PB文件已经存在则，返回PB文件的路径，否则将模型转化为PB文件，并且返回存储PB文件的路径
        if args.model_pb_dir is None:
            tmp_file = args.model_dir
        else:
            tmp_file = args.model_pb_dir

        pb_file = os.path.join(tmp_file, 'classification_model.pb')
        if os.path.exists(pb_file):
            print('pb_file exits', pb_file)
            return pb_file
        
        #增加 从label2id.pkl中读取num_labels, 这样也可以不用指定num_labels参数； 2019/4/17
        if not args.num_labels:
            num_labels, label2id, id2label = init_predict_var()
        #---

        graph = tf.Graph()
        with graph.as_default():
            with tf.Session() as sess:
                input_ids = tf.placeholder(tf.int32, (None, args.max_seq_len), 'input_ids')
                input_mask = tf.placeholder(tf.int32, (None, args.max_seq_len), 'input_mask')

                bert_config = modeling.BertConfig.from_json_file(os.path.join(args.bert_model_dir, 'bert_config.json'))

                loss, per_example_loss, logits, probabilities = create_classification_model(bert_config=bert_config, is_training=False,
                    input_ids=input_ids, input_mask=input_mask, segment_ids=None, labels=None, num_labels=num_labels)
                
                # pred_ids = tf.argmax(probabilities, axis=-1, output_type=tf.int32, name='pred_ids')
                # pred_ids = tf.identity(pred_ids, 'pred_ids')

                probabilities = tf.identity(probabilities, 'pred_prob')
                saver = tf.train.Saver()

            with tf.Session() as sess:
                sess.run(tf.global_variables_initializer())
                latest_checkpoint = tf.train.latest_checkpoint(args.model_dir)
                logger.info('loading... %s ' % latest_checkpoint )
                saver.restore(sess,latest_checkpoint )
                logger.info('freeze...')
                from tensorflow.python.framework import graph_util
                tmp_g = graph_util.convert_variables_to_constants(sess, graph.as_graph_def(), ['pred_prob'])
                logger.info('predict cut finished !!!')
        
        # 存储二进制模型到文件中
        logger.info('write graph to a tmp file: %s' % pb_file)
        with tf.gfile.GFile(pb_file, 'wb') as f:
            f.write(tmp_g.SerializeToString())
        return pb_file
    except Exception as e:
        logger.error('fail to optimize the graph! %s' % e, exc_info=True)




if __name__ == '__main__':
    pass
    
    parser = argparse.ArgumentParser(description='Trans ckpt file to .pb file')
    parser.add_argument('-bert_model_dir', type=str, required=True,
                        help='chinese google bert model path')
    parser.add_argument('-model_dir', type=str, required=True,
                        help='directory of a pretrained BERT model')
    parser.add_argument('-model_pb_dir', type=str, default=None,
                        help='directory of a pretrained BERT model,default = model_dir')
    parser.add_argument('-max_seq_len', type=int, default=128,
                        help='maximum length of a sequence,default:128')
    parser.add_argument('-num_labels', type=int, default=None,
                        help='length of all labels,default=2')
    parser.add_argument('-verbose', action='store_true', default=False,
                        help='turn on tensorflow logging for debug')
    args = parser.parse_args()
    optimize_class_model(args, logger=None)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187

2、服务端部署与启动

现在可以安装服务端了，使用的是 bert-base, 来自于项目BERT-BiLSTM-CRF-NER, 服务端只是该项目中的一个部分。 项目地址：https://github.com/macanv/BERT-BiLSTM-CRF-NER ，感谢Macanv同学提供这么好的项目。

这里要说明一下，我们经常会看到bert-as-service 这个项目的介绍，它只能加载BERT的预训练模型，输出文本向量化的结果。 而如果要加载fine-turing后的模型，就要用到 bert-base 了，详请请见：基于BERT预训练的中文命名实体识别TensorFlow实现

下载代码并安装 ：

pip install bert-base==0.0.7 -i https://pypi.python.org/simple
1

或者

git clone https://github.com/macanv/BERT-BiLSTM-CRF-NER
cd BERT-BiLSTM-CRF-NER/
python3 setup.py install
1
2
3

使用 bert-base 有三种运行模式，分别支持三种模型，使用参数-mode 来指定：

• NER 序列标注类型，比如命名实体识别；
• CLASS 分类模型，就是本文中使用的模型
• BERT 这个就是跟bert-as-service 一样的模式了

之所以要分成不同的运行模式，是因为不同模型对输入内容的预处理是不同的，命名实体识别NER是要进行序列标注； 而分类模型只要返回label就可以了。
安装完后运行服务，同时指定监听 HTTP 8091端口：

export BERT_BASE_DIR=/home/gildata/bert/vocab_file/chinese_L-12_H-768_A-12
export TRAINED_CLASSIFIER=/home/gildata/bert/bert/output

bert-base-serving-start  
-model_dir $TRAINED_CLASSIFIER 
-bert_model_dir $BERT_BASE_DIR 
-model_pb_dir $TRAINED_CLASSIFIER 
-mode CLASS 
-max_seq_len 128 
-http_port 8091 
-port 5575
-port_out 5576
1
2
3
4
5
6
7
8
9
10
11
12
13

注意：port 和 port_out 这两个参数是API调用的端口号， 默认是5555和5556,如果你准备部署多个模型服务实例，那一定要指定自己的端口号，避免冲突。 我这里是改为： 5575 和 5576

如果报错没运行起来，可能是有些模块没装上,都是 bert_base/server/http.py里引用的，装上就好了：

sudo pip install flask 
sudo pip install flask_compress
sudo pip install flask_cors
sudo pip install flask_json
1
2
3
4

运行服务后会自动生成很多临时的目录和文件，为了方便管理与启动，可建立一个工作目录，并把启动命令写成一个shell脚本。 这里创建的是service.sh，这样可以比较方便地设置服务器启动时自动启动服务，另外增加了每次启动时自动清除临时文件

#!/bin/bash
#chkconfig: 2345 80 90
#description: 启动BERT分类模型 

echo '正在启动 BERT SERVICE...'
cd /home/gildata/bert/service
sudo rm -rf tmp*

export BERT_BASE_DIR=/home/gildata/bert/vocab_file/chinese_L-12_H-768_A-12
export TRAINED_CLASSIFIER=/home/gildata/bert/bert/output

bert-base-serving-start  
-model_dir $TRAINED_CLASSIFIER 
-bert_model_dir $BERT_BASE_DIR 
-model_pb_dir $TRAINED_CLASSIFIER 
-mode CLASS 
-max_seq_len 128 
-http_port 8091 
-port 5575
-port_out 5576
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

补充说明一下内存的使用情况： BERT在训练时需要加载完整的模型数据，要用的内存是比较多的，差不多要10G，我这里用的是GTX 1080 Ti 11G。 但在训练完后，按上面的方式部署加载pb模型文件时，就不需要那么大了，上面也可以看到pb模型文件就是400多M。 其实只要你使用的是BERT base 预训练模型，最终的得到的pb文件大小都是差不多的。

3、端口测试

模型服务端部署完成了，可以使用curl命令来测试一下它的运行情况。

curl -X POST http://127.0.0.1:8091/encode
  -H 'content-type: application/json'
  -d '{"id": 111,"texts": ["总的来说，这款手机性价比是特别高的。","槽糕的售后服务！！！店大欺客"], "is_tokenized": false}'
1
2
3

执行结果：

>   -H 'content-type: application/json' \
>   -d '{"id": 111,"texts": ["总的来说，这款手机性价比是特别高的。","槽糕的售后服务！！！店大欺客"], "is_tokenized": false}'
{"id":111,"result":[{"pred_label":["1","-1"],"score":[0.9974544644355774,0.9961422085762024]}],"status":200}
1
2
3

可以看到对应的两个评论，预测结果一个是1，另一个是-1，计算的速度还是非常很快的。

因为现在很多的项目都是用java开发的，使用接口调用还是会有很多不便，下一章节，会介绍如何用java来调用bert模型。