docker部署tf-serving模型

自从使用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的接口就够用了。