自从使用docker制作python环境镜像之后,越来越觉得docker非常方便和友好,之前部署模型一直都是直接启用flask服务的形式,最近正好在弄bert模型,为了提高bert模型在cpu服务器上的推理效率,打算使用tf-serving的服务,虽然快不了多少,但是提升一点是一点,同时也省的每次直接部署的时候拷贝模型文件了。
一、模型准备
由于tf-serving需要模型静态图格式文件进行推理,因此首先将ckpt文件格式转成pb格式,这个步骤需要一个脚本文件,定义好输入输出,并加载预训练的ckpt文件,然后再保存成pb格式。
下面是参考官方源码修改的代码:
class BertServing(tf.keras.Model):
"""Bert transformer encoder model for serving."""
def __init__(self, config, bert_config, name_to_features, name="serving_model"):
super(BertServing, self).__init__(name=name)
cfg = bert_config
self.bert_encoder = BertEncoder(
vocab_size=cfg.vocab_size,
hidden_size=cfg.hidden_size,
num_layers=cfg.num_hidden_layers,
num_attention_heads=cfg.num_attention_heads,
intermediate_size=cfg.intermediate_size,
activation=tf_utils.get_activation(cfg.hidden_act),
dropout_rate=cfg.hidden_dropout_prob,
attention_dropout_rate=cfg.attention_probs_dropout_prob,
max_sequence_length=cfg.max_position_embeddings,
type_vocab_size=cfg.type_vocab_size,
initializer=tf.keras.initializers.TruncatedNormal(
stddev=cfg.initializer_range),
embedding_width=cfg.embedding_size,
return_all_encoder_outputs=True)
self.model = SentenceEmbedding(self.bert_encoder, config)
# ckpt = tf.train.Checkpoint(model=self.bert_encoder)
# init_checkpoint = self.config['bert_model_path']
#
# ckpt.restore(init_checkpoint).assert_existing_objects_matched()
self.name_to_features = name_to_features
def call(self, inputs):
input_word_ids = inputs["input_word_ids"]
input_mask = inputs["input_mask"]
input_type_ids = inputs["input_type_ids"]
infer_input = {
"input_word_ids": input_word_ids,
"input_mask": input_mask,
"input_type_ids": input_type_ids,
}
encoder_outputs = self.model(
infer_input)
return encoder_outputs
def serve_body(self, input_ids, input_mask=None, segment_ids=None):
if segment_ids is None:
# Requires CLS token is the first token of inputs.
segment_ids = tf.zeros_like(input_ids)
if input_mask is None:
# The mask has model1 for real tokens and 0 for padding tokens.
input_mask = tf.where(
tf.equal(input_ids, 0), tf.zeros_like(input_ids),
tf.ones_like(input_ids))
inputs = dict(
input_word_ids=input_ids, input_mask=input_mask, input_type_ids=segment_ids)
return self.call(inputs)
@tf.function
def serve(self, input_ids, input_mask=None, segment_ids=None):
outputs = self.serve_body(input_ids, input_mask, segment_ids)
# Returns a dictionary to control SignatureDef output signature.
return {"outputs": outputs}
@tf.function
def serve_examples(self, inputs):
features = tf.io.parse_example(inputs, self.name_to_features)
for key in list(features.keys()):
t = features[key]
if t.dtype == tf.int64:
t = tf.cast(t, tf.int32)
features[key] = t
return self.serve(
features["input_word_ids"],
input_mask=features["input_mask"] if "input_mask" in features else None,
segment_ids=features["input_type_ids"]
if "input_type_ids" in features else None)
@classmethod
def export(cls, model, export_dir):
if not isinstance(model, cls):
raise ValueError("Invalid model instance: %s, it should be a %s" %
(model, cls))
signatures = {
"serving_default":
model.serve.get_concrete_function(
input_ids=tf.TensorSpec(
shape=[None, None], dtype=tf.float32, name="inputs")),
}
if model.name_to_features:
signatures[
"serving_examples"] = model.serve_examples.get_concrete_function(
tf.TensorSpec(shape=[None], dtype=tf.string, name="examples"))
tf.saved_model.save(model, export_dir=export_dir, signatures=signatures)
def main(_):
config_path = FLAGS.config_path
with open(config_path, 'r') as fr:
config = json.load(fr)
sequence_length = config['seq_len']
if sequence_length is not None and sequence_length > 0:
name_to_features = {
"input_word_ids": tf.io.FixedLenFeature([sequence_length], tf.int64),
"input_mask": tf.io.FixedLenFeature([sequence_length], tf.int64),
"input_type_ids": tf.io.FixedLenFeature([sequence_length], tf.int64),
}
else:
name_to_features = None
bert_config = bert_configs.BertConfig.from_json_file(FLAGS.bert_config_file)
serving_model = BertServing(
config=config, bert_config=bert_config, name_to_features=name_to_features)
checkpoint = tf.train.Checkpoint(model=serving_model.bert_encoder)
checkpoint.restore(FLAGS.model_checkpoint_path
).assert_existing_objects_matched()
'''.run_restore_ops()'''
BertServing.export(serving_model, FLAGS.export_path)
我修改的地方主要是init模型,然后call、serve_body函数,改成自己的定义的模型推理就行,之后模型就会保存成pb格式,我的文件夹目录如下:
二、模型配置文件及多模型部署
在模型保存之前注意设置好模型的存储路径,需要注意的一点就是模型的版本,版本一般为数字,因为如果检测不到模型的版本,启动服务的时候会报错。参考上面图片,在model1路径下还有一个名为数字1的文件夹代表模型的版本。
如果要使用多个模型的服务,需要创建多个模型的路径,并编辑一个models.config文件,内容如下:
model_config_list:{
config:{
name:"model1",
base_path:"/models/my_model/model1"
model_platform:"tensorflow"
model_version_policy:{
all:{}
}
},
config:{
name:"model2",
base_path:"/models/my_model/model2"
model_platform:"tensorflow"
}
}
其中base_path为docker内部的文件路径,部署前需要将模型文件拷贝到docker内相应的路径下。
如果model1也有多个版本的模型路径,可以在config文件中添加:
model_version_policy:{ all:{} }
并且url要确定模型的版本,以下是几种url代表的含义:
选择版本:
'http://localhost:8501/v1/models/model1/versions/100001:predict'
默认最新版本:
'http://localhost:8501/v1/models/model1:predict'
选择model2:
'http://localhost:8501/v1/models/model2:predict'
三、接口调用
使用http调用的是8501端口号,如果使用grpc调用,端口号为8500。grpc接口多用于于图片数据调用,需要预先把数据转成tensor格式,因此接口速度会快一些。nlp用http的接口就够用了。