tf中checkpoint文件转saved_model，并通过tfserving请求

# ckpt转saved_model并自定义输入输出tensor名字
input_node = tf.placeholder(tf.float32, shape=(None, None, None, 3),name="input") 
logits_tf =  network.deeplab_v3(input_node, args, is_training=False, reuse=False)
print(logits_tf.shape)
saver = tf.train.Saver()
with tf.Session() as sess:

    # Create a saver.
    sess.run(tf.local_variables_initializer())
    sess.run(tf.global_variables_initializer())

    # Restore variables from disk.
    saver.restore(sess, checkpoint_file)
    builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
    signature = predict_signature_def(inputs={'input': input_node},
                                      outputs={'output': logits_tf})
    builder.add_meta_graph_and_variables(sess=sess,
                                         tags=[tag_constants.SERVING],
                                         signature_def_map={'predict': signature})
    builder.save() 

通过saved_model_cli显示的结果来看：

signature_def['predict']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['input'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, -1, -1, 3)
        name: input:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['output'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, -1, -1, 2)
        name: DeepLab_v3/ResizeBilinear:0
  Method name is: tensorflow/serving/predict

有一个坑，自定义输入输出tensor名字，是tf_serving使用时才有效的，而如果直接加载图的话，还是得取tensor原本的名字

from matplotlib import pyplot as plt
%matplotlib inline

import numpy as np
import cv2
import albumentations as A
with tf.Session(graph=tf.Graph()) as sess:
    tf.saved_model.loader.load(sess, ["serve"], export_dir)
    graph = tf.get_default_graph()

    img = cv2.imread("/home/ground_train/test/20112201814-1_00000036_img.png")
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = A.PadIfNeeded(p=1,min_height=1088, min_width=1920)(image=img)['image']
    img = img.astype(np.float32)
    img_x = np.expand_dims(img, 0)
    x = sess.graph.get_tensor_by_name('input:0')
    y = sess.graph.get_tensor_by_name('DeepLab_v3/ResizeBilinear:0')
    logits_tf = sess.run(y,feed_dict={x: img_x})
    print(logits_tf.shape)
    predictions_tf = np.argmax(logits_tf, axis=3)
    pred_image = predictions_tf[0]
    print(pred_image.shape)
    f, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 8))

    ax1.imshow(img.astype(np.uint8))
    ax2.imshow(pred_image)
    plt.show()

直接tf-serving请求：

启动：

docker run -p port:8500 -e MODEL_NAME=ground_segmentation --name tfserving-ground-segmentation -e NVIDIA_VISIBLE_DEVICES=1 -v /model_path_in_machine:/models/ground_segmentation docker_images:tag --per_process_gpu_memory_fraction=0.2

需要注意的是：/model_path_in_machine下应该是版本号，例如0，0下一级目录是saved_model

请求代码：

import grpc
import numpy as np
import tensorflow as tf
from PIL import Image
from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2_grpc
import cv2
import matplotlib.pyplot as plt
 
def request_tfserving(inputs, server_url, model_name, signature_name, input_names, output_names):
    # 建立连接
    channel = grpc.insecure_channel(server_url)
    stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
    # 开始设置请求数据
    request = predict_pb2.PredictRequest()
    request.model_spec.name = model_name  # 模型名称
    request.model_spec.signature_name = signature_name  # 签名名称（默认 serving_default）
    # 设置输入数据
    for input, input_name in zip(inputs, input_names):
        request.inputs[input_name].CopyFrom(tf.make_tensor_proto(input, shape=list(input.shape)))
        
    response = stub.Predict(request, 5.0)  # 其中第2个参数为请求的 timeout 时长
    res_from_server_np = []
    for output_name in output_names:
        res_from_server_np.append(tf.make_ndarray(response.outputs[output_name]))
    print(np.argmax(res_from_server_np[0][0,:],axis=2).shape)
    f, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 8))

    ax1.imshow(input[0,:].astype(np.uint8))
    ax2.imshow(np.argmax(res_from_server_np[0][0,:],axis=2))
    plt.show()
 
if __name__ == "__main__":
    img_path = '/home/ground_train/test/20112201814-1_00000036_img.png'
    img = Image.open(img_path)
    img = img.resize((960, 544), Image.NEAREST).convert('RGB')
    img = image.img_to_array(img)
    img = np.expand_dims(img, axis=0).astype('float32')
    print(img.shape)
    request_tfserving(inputs=[img],#这里一定得是个list
                      server_url='ip:port',
                      model_name='ground_segmentation',
                      signature_name='predict',#与saved_model_cli结果中第一行的key一致，默认是serving_default
                      input_names=['input'],
                      output_names=['output'])

 用restful API请求：

参照https://github.com/tensorflow/serving/blob/master/tensorflow_serving/g3doc/api_rest.md#example

docker run --rm -p port:8501 --mount type=bind,source=/home/work/maleiyuan/deeplab_v3,target=/home/work/maleiyuan/deeplab_v3 -e MODEL_BASE_PATH=/home/work/maleiyuan/deeplab_v3 -e MODEL_NAME=ground_segmentation -t docker_images:tag

其中，MODEL_BASE_PATH+MODEL_NAME是放模型文件的路径，其中MODEL_NAME是ground_segmentation(放模型文件的文件夹名字)

 curl http://localhost:port/v1/models/ground_segmentation，返回结果：

{
 "model_version_status": [
  {
   "version": "0",
   "state": "AVAILABLE",
   "status": {
    "error_code": "OK",
    "error_message": ""
   }
  }
 ]
}