NLP之BERT分类模型部署提供服务_label2id.pkl-优快云博客

本文详细介绍如何将BERT预训练模型微调后，转换为.pb格式，部署为服务端，包括生成label2id.pkl，使用freeze_graph.py转换模型，以及通过bert-base启动服务。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1. labek2id.pkl

2. graph.pb

3. 模型参数需要一致

在我们使用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、转换模型
在训练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]

这里要注意的参数是：

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实现

下载代码并安装：

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/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
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}'

执行结果：

> -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模型。
————————————————
版权声明：本文为优快云博主「Little Programmer」的原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.youkuaiyun.com/qq_20989105/article/details/93854316