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

今天药忘吃喽~ 2022-01-07 11:23 596阅读 0赞

在我们使用bert预分类模型微调之后（[可以参考我前面写的文章][Link 1]），需要对项目进行支持，那就需要分类模型落地提供服务，这篇文章介绍python调用bert模型，提供服务。  
参考：[https://github.com/xmxoxo/BERT-train2deploy][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、转换模型 ###

在训练bert模型之后会得到一个output文件夹，里面是tf的checkout文件，模型是.ckpt的文件格式,文件比较大，并且有多个文件：  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzIwOTg5MTA1_size_16_color_FFFFFF_t_70]

可以看到，模板文件非常大，大约有好几个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]

这里要注意的参数是：

*  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

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

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

附 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)

### 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实现][BERT_TensorFlow]

下载代码并安装 ：

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

或者

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

使用 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

注意：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

运行服务后会自动生成很多临时的目录和文件，为了方便管理与启动，可建立一个工作目录，并把启动命令写成一个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

补充说明一下内存的使用情况： 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}'

执行结果：

>   -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，另一个是-1，计算的速度还是非常很快的。

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

[Link 1]: https://blog.csdn.net/qq_20989105/article/details/89492442
[https_github.com_xmxoxo_BERT-train2deploy]: https://github.com/xmxoxo/BERT-train2deploy
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzIwOTg5MTA1_size_16_color_FFFFFF_t_70]: /images/20211223/5660ee377c074b1c90c34b288984cda3.png
[BERT_TensorFlow]: https://blog.csdn.net/macanv/article/details/85684284