Tensorflow Serving 生产环境部署介绍

Tensorflow Serving 生产环境部署介绍

TF-serving介绍

TensorFlow Serving是google提供的一种生产环境部署方案,一般来说在做算法训练后,都会导出一个模型,在应用中直接使用。

正常的思路是在flask或者tornado这种web服务中嵌入tensorflow的模型,提供rest api的云服务接口。考虑到并发高可用性,一般会采取多进程的部署方式,即一台云服务器上同时部署多个flask,每个进程独享一部分GPU资源,显然这样是很浪费资源的。

Google提供了一种生产环境的新思路,他们开发了一个tensorflow-serving的服务,可以自动加载某个路径下的所有模型,模型通过事先定义的输入输出和计算图,直接提供rpc或者rest的服务。

  • 一方面,支持多版本的热部署(比如当前生产环境部署的是1版本的模型,训练完成后生成一个2版本的模型,tensorflow会自动加载这个模型,停掉之前的模型)。
  • 另一方面,tensorflow serving内部通过异步调用的方式,实现高可用,并且自动组织输入以批次调用的方式节省GPU计算资源。

因此,整个模型的调用方式就变成了:

客户端 ----> web服务 --grpc或者rest--> tensorflow serving

如果我们想要替换模型或者更新版本,只需要训练模型并将训练结果保存到固定的目录下就可以了。

环境准备

  • 首先需要安装nvidia-driver(gpu驱动)以及Docker 19.03
  • 安装nvidia-docker,这是nvidia在docker上进行了封装,让docker可以使用GPU资源,具体安装方法可以参考以下链接:github.com/NVIDIA/nvidi

安装命令如下相关命令如下:

shell   distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.repo | sudo tee /etc/yum.repos.d/nvidia-docker.repo
sudo yum install -y nvidia-container-toolkit
sudo systemctl restart docker
  • 拉取TFserving的GPU镜像

shell docker pull tensorflow/serving-gpu

制作模型文件

低阶API版本

TF-serving需要使用的模型是pb模型文件,而不是通常使用的ckpt模型文件,因此需要指定相应的参数。

以下是一个可以用来生成pb模型的代码,参考mnist_saved_model.py

  • 确定好模型的输出路径,模型的输入路径是一个由一串数字命名的文件夹,数字就是版本号
output_dir = "counter"
if not os.path.exists(output_dir):
    os.mkdir(output_dir)
for i in range(100000, 9999999):
    cur = os.path.join(output_dir, str(i))
    if not tf.gfile.Exists(cur):
        output_dir = cur
        break
  • 建立一个模型生成的工具类SavedModelBuilder:
builder = tf.saved_model.builder.SavedModelBuilder(output_dir)
  • 构建模型可供调用的方法签名,以及输入和输出的类型,其中get_counter,incr_counter,incr_counter_by,reset_counter分别对应着四个方法的签名
signature_def_map = signature({
        "get_counter": {"inputs": {"nothing": nothing},
                        "outputs": {"output": counter},
                        "method_name": method_name},
        "incr_counter": {"inputs": {"nothing": nothing},
                         "outputs": {"output": incr_counter},
                         "method_name": method_name},
        "incr_counter_by": {"inputs": {'delta': delta, },
                            "outputs": {'output': incr_counter_by},
                            "method_name": method_name},
        "reset_counter": {"inputs": {"nothing": nothing},
                          "outputs": {"output": reset_counter},
                          "method_name": method_name}
    })
  • 而在进行定义输入和输出的类型的时候,我们抽出了一个函数,协助处理,其中tf.saved_model.utils.build_tensor_info可以根据传入的tensor 对象构建protocol buffer,每个方法签名都会构建一个对象,并最终生成一个signature_dict,在后续请求方法的时候,request.model_spec.signature_name需要制定这些方法签名的key值。同时我们注意到这里的input里面还有一个key值,这个key是request.inputs['nothing']这里制定。还有一个参数method_name,是用来表示该方法属于预测,还是分类或者回归。
def signature(function_dict):
    signature_dict = {}
    for k, v in function_dict.items():
        inputs = {k: tf.saved_model.utils.build_tensor_info(v) for k, v in v['inputs'].items()}
        outputs = {k: tf.saved_model.utils.build_tensor_info(v) for k, v in v['outputs'].items()}
        signature_dict[k] = tf.saved_model.build_signature_def(inputs=inputs, outputs=outputs,
                                                               method_name=v['method_name'])
    return signature_dict
  • 添加需要存储的信息,其中tag要使用[tf.compat.v1.saved_model.tag_constants.SERVING]表明是要提供给serving的。main_op=tf.tables_initializer(), strip_default_attrs=True这两个参数是用来初始化一个lookup_table以及版本兼容用的
builder.add_meta_graph_and_variables(sess, tags=[tf.saved_model.tag_constants.SERVING],
signature_def_map=signature_def_map, main_op=tf.tables_initializer(), strip_default_attrs=True)
  • 保存模型
builder.save()

完整的代码如下:

from __future__ import division, absolute_import, print_function

import os

import tensorflow.compat.v1 as tf  # tf2.1兼容

tf.disable_v2_behavior()


def signature(function_dict):
    signature_dict = {}
    for k, v in function_dict.items():
        inputs = {k: tf.saved_model.utils.build_tensor_info(v) for k, v in v['inputs'].items()}
        outputs = {k: tf.saved_model.utils.build_tensor_info(v) for k, v in v['outputs'].items()}
        signature_dict[k] = tf.saved_model.build_signature_def(inputs=inputs, outputs=outputs,
                                                               method_name=v['method_name'])
    return signature_dict


output_dir = "counter"
if not os.path.exists(output_dir):
    os.mkdir(output_dir)
for i in range(100000, 9999999):
    cur = os.path.join(output_dir, str(i))
    if not tf.gfile.Exists(cur):
        output_dir = cur
        break
method_name = tf.saved_model.signature_constants.PREDICT_METHOD_NAME
builder = tf.saved_model.builder.SavedModelBuilder(output_dir)
print('outputdir', output_dir)
with tf.Graph().as_default(), tf.Session() as sess:
    counter = tf.Variable(0.0, dtype=tf.float32, name="counter")
    with tf.name_scope("incr_counter_op", values=[counter]):
        incr_counter = counter.assign_add(1.0)
    delta = tf.placeholder(dtype=tf.float32, name="delta")
    with tf.name_scope("incr_counter_by_op", values=[counter, delta]):
        incr_counter_by = counter.assign_add(delta)
    with tf.name_scope("reset_counter_op", values=[counter]):
        reset_counter = counter.assign(0.0)
    nothing = tf.placeholder(dtype=tf.int32, shape=(None,))
    sess.run(tf.global_variables_initializer())
    signature_def_map = signature({
        "get_counter": {"inputs": {"nothing": nothing},
                        "outputs": {"output": counter},
                        "method_name": method_name},
        "incr_counter": {"inputs": {"nothing": nothing},
                         "outputs": {"output": incr_counter},
                         "method_name": method_name},
        "incr_counter_by": {"inputs": {'delta': delta, },
                            "outputs": {'output': incr_counter_by},
                            "method_name": method_name},
        "reset_counter": {"inputs": {"nothing": nothing},
                          "outputs": {"output": reset_counter},
                          "method_name": method_name}
    })

    builder.add_meta_graph_and_variables(sess, tags=[tf.saved_model.tag_constants.SERVING],
                                         signature_def_map=signature_def_map, main_op=tf.tables_initializer(),
                                         strip_default_attrs=True)
    builder.save()
    print("over")

tf.estimator版本

如果我们使用tf.estimator导出的话,也需要提供输入和输出,输出需要在模型的预测代码返回的实例里面指出,可参考如下代码,再返回预测的实例对象时,传入了export_outputs参数。

if mode == tf.estimator.ModeKeys.PREDICT:
    predictions = {
      'probabilities': prop,
      'ctr_probabilities': ctr_predictions,
      'cvr_probabilities': cvr_predictions
    }
    export_outputs = {
      'prediction': tf.estimator.export.PredictOutput(predictions)
    }
    return tf.estimator.EstimatorSpec(mode, predictions=predictions, export_outputs=export_outputs)

而输入则是需要在主动保存模型时添加的,首先要构建一个serving_input_receiver_fn,用来告诉模型应该接受什么样的输入,这里的receiver_tensors就是最后tensorflow serving需要接受的参数。

官方建议使用传入tf.example对象,然后再解析成为tensor,但还是有点麻烦,因为客户端也得传入这个对象才可以。

feature_spec = {'foo': tf.FixedLenFeature(...),
                'bar': tf.VarLenFeature(...)}

def serving_input_receiver_fn():
  """An input receiver that expects a serialized tf.Example."""
  serialized_tf_example = tf.placeholder(dtype=tf.string,shape=[default_batch_size],name='input_example_tensor')
  receiver_tensors = {'examples': serialized_tf_example}
  features = tf.parse_example(serialized_tf_example, feature_spec)
  return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)

但是如果我们在客户端想直接传入tensor,那么可以作如下改写:

def serving_input_receiver_fn():

    tensor1 = tf.placeholder(dtype=tf.int32,shape=[None,20],name='tensor1')
    tensor2 = tf.placeholder(dtype=tf.int32,shape=[None,10],name='tensor2')
    receiver_tensors = {'tensor1': tensor1,'tensor2': tensor2}
    features = receiver_tensors
    return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)
解释一下这里的receiver_tensors和features的区别,receiver_tensors是从客户端的request接收到的tensor,对应的是placeholder信息,而features则是指输入到model_fn的feature.假如我们接受的是tf.example对象,那么我们需要先将他parse变成相应的tensor。但是这里我们不需要进行这个操作,因为接收到的请求参数已经是可以直接扔到model_fn的tensor了,不用转换。可参考链接里的解释:TensorFlow Estimator ServingInputReceiver features vs receiver_tensors: when and why?

最后使用estimator实例的export_savedmodel方法导出模型到export_dir_base文件夹

estimator.export_savedmodel(export_dir_base,serving_input_receiver_fn,strip_default_attrs=True)

tf.keras版本

参考链接:使用REST训练和提供模型,参数签名根据输入和输出的默认张量名进行匹配

tf.keras.models.save_model(
    model,
    export_path,
    overwrite=True,
    include_optimizer=True,
    save_format=None,
    signatures=None,
    options=None
)

# 或者
model.save("my_model_dir",save_format='tf')

tensorflow serving部署模型

首先我们有如下的文件结构

tmp
└── counter
    └── 100000
    │  ├── saved_model.pb
    │  └── variables
    │    ├── variables.data-00000-of-00001
    │    └── variables.index
    └── 100001
       ├── saved_model.pb
       └── variables
         ├── variables.data-00000-of-00001
         └── variables.index

那我们可以用如下的命令去启动该模型,其中8500是gRPC端口,8501 是 REST API的端口,也可以只开启其中一个。可用通过-e NVIDIA_VISIBLE_DEVICES=0参数指定哪块GPU去运行程序

docker run --runtime=nvidia -p 8500:8500 -p 8501:8501 --mount type=bind,source=/tmp/counter,target=/models/counter -e MODEL_NAME=counter -t tensorflow/serving-gpu &

如果需要部署多个模型,那模型的文件需要用如下的组织形式

multiModel
├── counter
│   └── 100000
│      ├── saved_model.pb
│      └── variables
│        ├── variables.data-00000-of-00001
│        └── variables.index
├── counter1
│     └── 100000
│        ├── saved_model.pb
│        └── variables
│          ├── variables.data-00000-of-00001
│          └── variables.index
└── models.config

其中需要包含一个models.config文件,该文件会告知需要部署哪些模型。文件是ASCII protocol buffer的结构,具体什么是ASCII protocol buffer,可参考链接

model_config_list:{
    config:{
      name:"counter",
      base_path:"/models/multiModel/counter",
      model_platform:"tensorflow",
      model_version_policy:{
        # 这是加载全部模型的策略
        all:{}
      }
      version_labels:{
          key:"stable",
          value:100000
      }
    },
    config:{
      name:"counter1",
      base_path:"/models/multiModel/counter1",
      model_platform:"tensorflow",
      model_version_policy:{
        # 这是指定加载version的策略
        specific:{
            version:100000
        }
      }
    },
}

启动TFserving的服务类似于下面的命令。这里需要注意的是allow_version_labels_for_unavailable_models参数需要传个true进去,因为我们之前在model.config里面定义的模型是all:{}策略,没有指定加载模型version,不指定这个参数,启动容器会报错。还有一个参数--model_config_file_poll_wait_seconds=60,这个参数可以定期检查config文件,然后动态改变serve的模型,这两个参数需要放在最后。

docker run --runtime=nvidia -p 8500:8500 -p 8501:8501 --name tf_serving --mount type=bind,source=/home/node1/model/multiModel/,target=/models/multiModel -t tensorflow/serving:latest-gpu --model_config_file=/models/multiModel/models.config --model_config_file_poll_wait_seconds=60 --allow_version_labels_for_unavailable_models=true

请求tensorflow serving的预测服务

我们可以使用如下命令去获取serve的模型的方法和参数签名信息

crul http://host:8501/v1/models/${MODEL_NAME}[/versions/${MODEL_VERSION}]/metadata

REST API接口请求方式

(请求参数和上面的模型没关系,只是一个例子)

参考链接serving/api_rest

如果使用rest ful形式的去请求服务,请求的url类似如下

POST http://host:8501/v1/models/${MODEL_NAME}[/versions/${MODEL_VERSION}]:predict

其中/versions/${MODEL_VERSION}是可选的,如果不添加,则表示使用最新版本的模型,也就是MODEL_VERSION最大的那个模型。

请求的body是一个json字符串,body有两种模式,行模式(或者叫instance模式)和列模式

行模式如下:

{
  "signature_name": <string>,
  "instances": <value>|<(nested)list>|<list-of-objects>
}

列模式如下:

{
  "signature_name": <string>,
  "inputs": <value>|<(nested)list>|<object>
}

其中"signature_name"字段表示的是模型的方法签名,也就是之前定义的signature_def_map里面的值,默认应该是 tf.compat.v1.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY,也就是“serving_default"

对于行模式,加入需要传入一个实例为

{
     "tag": ["foo"],
     "signal": [1, 2, 3, 4, 5],
     "sensor": [1, 2, 3, 4]
   }

也就是传入"tag","signal","sensor"这几个key对应的数值,这三个值应当与之前prediction_signature里面定义的input的key对应,并且注意的这三个命名变量一定要拥有相同的0维,如果不是,则需要使用列模式。

这里我使用的例子和官方的有区别,因为官方的描述是有矛盾的,参考了Tensorflow Serving-Docker RESTful API客户端访问问题排查里面的案例描述,暂时认为官方文字描述是正确的。事实上,如果每个元素的0维不一致,我们也可以在模型输入的外层套一个维度,也可以满足0维相同的要求。

Note, each named input ("tag", "signal", "sensor") is implicitly assumed have same 0-th dimension (two in above example, as there are two objects in the instances list). If you have named inputs that have different 0-th dimension, use the columnar format described below.

如果需要传入多个值

{
 "instances": [
   {
     "tag": ["foo"],
     "signal": [1, 2, 3, 4, 5],
     "sensor": [1, 2, 3, 4]
   },
   {
     "tag": ["bar"],
     "signal": [3, 4, 1, 2, 5],
     "sensor": [4, 5, 6, 8]
   }
 ]
}

行模式的返回值是一个json字符串,如果模型的输出只包含一个命名的tensor,我们省略名字,直接使用标量或者值的list。如果模型输出多个命名的tensor,我们输出对象list,和上面提到的行形式输入类似。

{
  "predictions": <value>|<(nested)list>|<list-of-objects>
}

列模式

{
 "inputs": {
   "tag": ["foo", "bar"],
   "signal": [[1, 2, 3, 4, 5], [3, 4, 1, 2, 5]],
   "sensor": [[[1, 2], [3, 4]], [[4, 5], [6, 8]]]
 }
}

可以看出这里的每个key值后面跟着多个Tensor,他们是一一对应的,而且这里并不需要每个元素具有相同的0维

列模式的返回值也是json字符串,key是outputs,如果模型的输出只包含一个命名的tensor,我们省略名字和outputs key map,直接使用标量或者值的list。如果模型输出多个命名的tensor,我们输出对象,其每个key都和输出的tensor名对应,和上面提到的列形式输入类似。

{
  "outputs": <value>|<(nested)list>|<object>
}

Grpc请求方式

python版本Grpc调用

Grpc调用是需要proto文件来生成一些依赖代码,相关proto文件在链接里面。

version1:使用封装好的工具进行调用

编译proto文件这一步,显然有人会已经帮我们做好了,并打包上传了名为tensorflow-serving-api的工具。我们可以从中直接获取对应的依赖文件。以下是一个利用它建立依赖的过程,首先需要指定model_name和signature_name,model_name是模型签名,signature_name模型方法签名。传入的tensor需要经过tf.contrib包进行转换成protobuf形式,tensor的key则设置为存储模型的时候指定的key值。

注意:这里如果想要指定模型的version,有两种办法。一:如果是在model.config文件里写入了version_label,那么就可以用request.model_spec.version_label='stable'这种办法指定;二:如果没有,则可以使用version数字,传入request.model_spec.version.value=00000123之类的版本号即可。在这里的proto文件里面,使用了oneof语法,也就是只会接受一种,如果同时传入了两中形式,那么会使用后写入的版本

  request = predict_pb2.PredictRequest()
  request.model_spec.name = 'counter'
  request.model_spec.version_label='stable'
  request.model_spec.version.value=00000123  # 同时传入version_label和version的话,只有写在后面的代码会生效
  request.model_spec.signature_name = 'incr_counter'

  # read image into numpy array
  inputs=np.array([0])

  # convert to tensor proto and make request
  # shape is in NHWC (num_samples x height x width x channels) format
  tensor = tf.contrib.util.make_tensor_proto(inputs, shape=list(inputs.shape))
  request.inputs['nothing'].CopyFrom(tensor)

完整的代码如下

from __future__ import print_function
import numpy as np
import time
tt = time.time()

import tensorflow as tf  # tf1.x,需要使用里面的contrib包,tf2.x里面没有了

from grpc
from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2_grpc


def main():  
  # create prediction service client stub
  channel = grpc.insecure_channel("172.0.0.1:8501")
  stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)

  # create request
  request = predict_pb2.PredictRequest()
  request.model_spec.name = 'counter'
  request.model_spec.version_label='stable'
  request.model_spec.version.value=100000  # 同时传入version_label和version的话,只有写在后面的代码会生效
  request.model_spec.signature_name = 'incr_counter'
  input = np.array([0])

  # convert to tensor proto and make request
  # shape is in NHWC (num_samples x height x width x channels) format
  tensor = tf.contrib.util.make_tensor_proto(input, shape=list(input.shape))
  request.inputs['nothing'].CopyFrom(tensor)
  resp = stub.Predict(request, 30.0)

  print('total time: {}s'.format(time.time() - tt))

if __name__ == '__main__':
    main()

version2:自己编译proto文件生成依赖

上述两个模块需要太多的依赖,而实际上我们并不需要这么多依赖,因此可以使用自己编译的proto文件,生成需要的依赖

参考链接一 参考链接二

首先从tensorflow和tensorflow serving的github里面下载proto文件

tensorflow/serving/  
  tensorflow_serving/apis/model.proto
  tensorflow_serving/apis/predict.proto
  tensorflow_serving/apis/prediction_service.proto

tensorflow/tensorflow/  
  tensorflow/core/framework/resource_handle.proto
  tensorflow/core/framework/tensor_shape.proto
  tensorflow/core/framework/tensor.proto
  tensorflow/core/framework/types.proto

将上述文件保存至protos文件

protos/  
  tensorflow_serving/
    apis/
      *.proto
  tensorflow/
    core/
      framework/
        *.proto

为了简单起见,prediction_service.proto(预测服务)可以简化为只实现Predict RPC。这避免了引入服务中定义的其他RPC的嵌套依赖关系。

使用grpcio.tools.protoc去编译生成依赖

PROTOC_OUT=protos/  
PROTOS=$(find . | grep "\.proto$")  
for p in $PROTOS; do  
  python -m grpc.tools.protoc -I . --python_out=$PROTOC_OUT --grpc_python_out=$PROTOC_OUT $p
done

然后就可以去除掉tensorflow-serving-api的依赖,同时我们可以用tensorflow里面proto生成的依赖文件,从而去除掉tensorflow的依赖,一般正常使用会用到tf.contrib.util.make_tensor_proto函数去根据numpy数组生成protocol buff,不需要引入这个依赖

我这里已经生成了一份

链接:pan.baidu.com/s/1ZcJplX 提取码:aomd

from __future__ import print_function, division, absolute_import

import time

import numpy as np

tt = time.time()

import grpc
from protos.tensorflow_serving.apis import predict_pb2
from protos.tensorflow_serving.apis import prediction_service_pb2_grpc
from protos.tensorflow.core.framework import tensor_pb2
from protos.tensorflow.core.framework import tensor_shape_pb2
from protos.tensorflow.core.framework import types_pb2


def incr_counter():
    # create prediction service client stub
    channel = grpc.insecure_channel("172.0.0.1:8501")
    stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
    # # create request
    request = predict_pb2.PredictRequest()
    request.model_spec.name = 'counter'
    request.model_spec.signature_name = 'incr_counter'

    input = np.array([0])
    tensor_shape = list(input.shape)
    dims = [tensor_shape_pb2.TensorShapeProto.Dim(size=dim) for dim in tensor_shape]
    print("+++++++++++")
    print(dims)
    tensor_shape = tensor_shape_pb2.TensorShapeProto(dim=dims)
    tensor = tensor_pb2.TensorProto(
        dtype=types_pb2.DT_INT32,
        tensor_shape=tensor_shape,
        int_val=list(input.reshape(-1)))
    print("+++++++++++")
    print(tensor)
    request.inputs['nothing'].CopyFrom(tensor)
    resp = stub.Predict(request, 5.0)
    print(resp)
    print('total time: {}s'.format(time.time() - tt))


def get_counter():
    # create prediction service client stub
    channel = grpc.insecure_channel("172.0.0.1:8501")
    stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
    #
    # # create request
    request = predict_pb2.PredictRequest()
    request.model_spec.name = 'counter'
    request.model_spec.signature_name = 'get_counter'

    input = np.array([0])
    tensor_shape = list(input.shape)
    dims = [tensor_shape_pb2.TensorShapeProto.Dim(size=dim) for dim in tensor_shape]
    print("+++++++++++")
    print(dims)
    tensor_shape = tensor_shape_pb2.TensorShapeProto(dim=dims)
    tensor = tensor_pb2.TensorProto(
        dtype=types_pb2.DT_INT32,
        tensor_shape=tensor_shape,
        int_val=list(input.reshape(-1)))
    print("+++++++++++")
    print(tensor)
    request.inputs['nothing'].CopyFrom(tensor)
    resp = stub.Predict(request, 5.0)
    print(resp)
    print('total time: {}s'.format(time.time() - tt))


def incr_counter_by():
    # create prediction service client stub
    channel = grpc.insecure_channel("172.0.0.1:8501")
    stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
    #
    # # create request
    request = predict_pb2.PredictRequest()
    request.model_spec.name = 'counter'
    request.model_spec.signature_name = 'incr_counter_by'

    input = 2
    # 这里需要输入的是一个scalar,不能有任何维度
    tensor_shape = tensor_shape_pb2.TensorShapeProto(dim=[])
    tensor = tensor_pb2.TensorProto(
        dtype=types_pb2.DT_FLOAT,
        tensor_shape=tensor_shape,
        float_val=[input])
    print("+++++++++++")
    print(tensor)
    request.inputs['delta'].CopyFrom(tensor)
    resp = stub.Predict(request, 5.0)
    print(resp)
    print('total time: {}s'.format(time.time() - tt))


def reset_counter():
    # create prediction service client stub
    channel = grpc.insecure_channel("172.0.0.1:8501")
    stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)

    # # create request
    request = predict_pb2.PredictRequest()
    request.model_spec.name = 'counter'
    request.model_spec.signature_name = 'get_counter'

    input = np.array([0])
    tensor_shape = list(input.shape)
    dims = [tensor_shape_pb2.TensorShapeProto.Dim(size=dim) for dim in tensor_shape]
    print("+++++++++++")
    print(dims)
    tensor_shape = tensor_shape_pb2.TensorShapeProto(dim=dims)
    tensor = tensor_pb2.TensorProto(
        dtype=types_pb2.DT_INT32,
        tensor_shape=tensor_shape,
        int_val=list(input.reshape(-1)))
    print("+++++++++++")
    print(tensor)
    request.inputs['nothing'].CopyFrom(tensor)
    resp = stub.Predict(request, 5.0)
    print(resp)
    print('total time: {}s'.format(time.time() - tt))


if __name__ == '__main__':
    incr_counter()
    get_counter()
    incr_counter_by()
    reset_counter()

这里解释一下根据proto文件构建协议体的过程

input = np.array([0])
tensor_shape = list(input.shape)
# 首先根据输入的numpy数组的shape构建protobuf维度信息,其最终在metadata里面的结构是dim:[{size:1,name:""},{size:2,name:""}]这样的
dims = [tensor_shape_pb2.TensorShapeProto.Dim(size=dim) for dim in tensor_shape]
print("+++++++++++")
print(dims)
tensor_shape = tensor_shape_pb2.TensorShapeProto(dim=dims)
# 然后构建tensor的protobuf,dtype有许多类型,可以依据所需选取,需要注意的是,如果传入的是DT_INT32之类的,那么传入实际值需要使用int_val参数,而不是float_val之类的,否则服务不会报错,但是结果是不对滴。具体需要传入什么值,可以依据tensor.proto文件里面的说明。
tensor = tensor_pb2.TensorProto(
    dtype=types_pb2.DT_INT32,
    tensor_shape=tensor_shape,
    int_val=list(input.reshape(-1)))

由于我们接受的返回值也是proto协议约定的,所以要获取numpy数组,需要进行解析,里面涉及的数据类型按照模型的实际值修改即可。

result_dict = dict()
for key in resp.outputs:
    tensor_proto = resp.outputs[key]
    shape = [d.size for d in tensor_proto.tensor_shape.dim]
    values = np.fromiter(tensor_proto.float_val, dtype=np.float)
    result_dict[key] = values.reshape(shape)

Java版本Grpc调用

java版本调用这里也是需要先试用proto文件生成对应的依赖,proto文件结构和之前的一样。然后根据依赖编写grpc客户端。这是一个对应counter模型的测试案例。

package com.meituan.test;

import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;

import java.util.Arrays;
import java.util.List;

import org.tensorflow.framework.DataType;
import org.tensorflow.framework.TensorProto;
import org.tensorflow.framework.TensorShapeProto;

import tensorflow.serving.Model;
import tensorflow.serving.Predict;
import tensorflow.serving.PredictionServiceGrpc;
public class App {
    public static void main(String[] args) {
        List<Integer> intList =Arrays.asList(1);
        ManagedChannel channel = ManagedChannelBuilder.forAddress("0.0.0.0", 8500).usePlaintext(true).build();
        //这里还是先用block模式
        PredictionServiceGrpc.PredictionServiceBlockingStub stub = PredictionServiceGrpc.newBlockingStub(channel);
        //创建请求
        Predict.PredictRequest.Builder predictRequestBuilder = Predict.PredictRequest.newBuilder();
        //模型名称和模型方法名预设
        Model.ModelSpec.Builder modelSpecBuilder = Model.ModelSpec.newBuilder();
        modelSpecBuilder.setName("counter");
        modelSpecBuilder.setSignatureName("incr_counter");
        modelSpecBuilder.setVersion(Int64Value.newBuilder().setValue(100000).build());
        predictRequestBuilder.setModelSpec(modelSpecBuilder);
        //设置入参,访问默认是最新版本,如果需要特定版本可以使用tensorProtoBuilder.setVersionNumber方法
        TensorProto.Builder tensorProtoBuilder = TensorProto.newBuilder();
        tensorProtoBuilder.setDtype(DataType.DT_INT32);
        TensorShapeProto.Builder tensorShapeBuilder = TensorShapeProto.newBuilder();

        tensorShapeBuilder.addDim(TensorShapeProto.Dim.newBuilder().setSize(1));
        tensorProtoBuilder.setTensorShape(tensorShapeBuilder.build());
        tensorProtoBuilder.addAllIntVal(intList);
        predictRequestBuilder.putInputs("nothing", tensorProtoBuilder.build());
        //访问并获取结果
        Predict.PredictResponse predictResponse = stub.predict(predictRequestBuilder.build());
        org.tensorflow.framework.TensorProto result=predictResponse.toBuilder().getOutputsOrThrow("output");
        System.out.println("预测值是:"+result.getFloatValList());
    }
}

简书主页:

612twilight - 简书www.jianshu.com图标

编辑于 03-16