用docker的方式发布tensorflow serving

author:hjf
date:2018-12-12

其实可以直接看tensorflow官方的教程，传送门|￣|

一、安装docker

|￣|

二、用docker拉取一个已经集成了tensorflow serving所需环境的image（镜像）

$ docker pull tensorflow/serving

三、把你的模型整成savedModel格式

Train and export TensorFlow model
savedModel README

示例代码：

# coding: utf-8
from __future__ import print_function
import tensorflow as tf
from nets.resnet_v1 import resnetv1

old_model_path = '/media/coolman/DATA/02代码/local/12wheel_detect/tf-faster-rcnn/output/res101/voc_Mixed_trainval/model/v2-20181113/res101_faster_rcnn_iter_430000.ckpt'
export_path='/media/coolman/DATA/02代码/local/12wheel_detect/tf-faster-rcnn/output/res101/voc_Mixed_trainval/saved_model/2'

CLASSES=('__background__',
           'level_0', 'level_1', 'level_2', 'level_3',
           'level_4', 'level_5', 'level_6', 'level_7', 'level_8', 'level_9')

# 1、构建网络结构，整成一个对象的形式
model = resnetv1(num_layers=101)
model.create_architecture("TEST", len(CLASSES),
                          tag='default', anchor_scales=[8, 16, 32])

def build_and_saved_wdl():
    # 2、启动一个session，并用已保存的模型结构和参数初始化一个saver
    session = tf.Session()
    session.run(tf.global_variables_initializer())
    saver = tf.train.Saver()
    saver.restore(sess=session, save_path=old_model_path)

    # 3、创建一个SavedModelBuilder
    builder = tf.saved_model.builder.SavedModelBuilder(export_path)
    # 4、定义输入
    inputs = {
        # 注意，这里是你预测模型的时候需要传的参数，调用模型的时候，传参必须和这里一致
        # 这里的model.input_x和model.keep_prob就是模型里面定义的输入placeholder
        "input_image": tf.saved_model.utils.build_tensor_info(model._image),
        "input_im_info": tf.saved_model.utils.build_tensor_info(model._im_info)
    }
    # 5、定义输出
    output = {"output_cls_score": tf.saved_model.utils.build_tensor_info(model._predictions["cls_score"]),
              "output_rois": tf.saved_model.utils.build_tensor_info(model._predictions["rois"]),
              "output_cls_prob": tf.saved_model.utils.build_tensor_info(model._predictions["cls_prob"]),
              "output_bbox_pred": tf.saved_model.utils.build_tensor_info(model._predictions["bbox_pred"])}
    # ６、定义SavedMedel的关键属性
    prediction_signature = tf.saved_model.signature_def_utils.build_signature_def(
        inputs=inputs,
        outputs=output,
        method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME
    )
    builder.add_meta_graph_and_variables(
        session,
        [tf.saved_model.tag_constants.SERVING],
        {tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: prediction_signature}
    )
    # 7、保存SavedModel
    builder.save()

if __name__ == '__main__':
    build_and_saved_wdl()

四、使用 CLI 检查并执行 SavedModel

我没试过，传送门|￣|

五、模型发布

docker run -p 8501:8501 \
  --mount type=bind,source=/path/to/my_model/,target=/models/my_model \
  -e MODEL_NAME=my_model \
  -t tensorflow/serving

run:启动一个container
-p：本机端口与container端口映射关系，前面的代表本机端口
–mount：主机上的文件目录挂载到container的某个目录下。type=bind表示绑定挂载，source代表主机目录，target代表container中的目录（固定结构/models/）。
-e： 传入环境变量（我猜的）
-t：指定image

六、启动能够使用GPU资源的container

懒得写了，Serving with Docker using your GPU