TensorFlow Serving介绍

简介

TensorFlow的模型文件包含了深度学习模型的Graph和所有参数，其实就是checkpoint文件，用户可以加载模型文件继续训练或者对外提供Inference服务。

使用SavedModel导出模型

模型导出方式参考 https://tensorflow.github.io/serving/serving_basic 。

使用方法基本如下。

  1. from tensorflow.python.saved_model import builder as saved_model_builder
  3. export_path_base = sys.argv[-1]
  4. export_path = os.path.join(
  5.       compat.as_bytes(export_path_base),
  6.       compat.as_bytes(str(FLAGS.model_version)))
  7. print 'Exporting trained model to', export_path
  9. builder = saved_model_builder.SavedModelBuilder(export_path)
  10. builder.add_meta_graph_and_variables(
  11.       sess, [tag_constants.SERVING],
  12.       signature_def_map={
  13.            'predict_images':
  14.                prediction_signature,
  15.            signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
  16.                classification_signature,
  17.       },
  18.       legacy_init_op=legacy_init_op)
  20. builder.save()

可以参考 https://github.com/tobegit3hub/deep_recommend_system/ 提供的可运行代码示例。

  1. ./dense_classifier.py --mode savedmodel

使用exporter导出模型

这里有导出TensorFlow serving支持的模型文件例子，可以参考使用 https://github.com/tobegit3hub/deep_recommend_system/blob/master/dense_classifier.py 。

导出的代码也比较简单，用户在inputs和output中填入模型Inference时的输入和输出即可。

  1. from tensorflow.contrib.session_bundle import exporter
  3. flags = tf.app.flags
  4. FLAGS = flags.FLAGS
  5. flags.DEFINE_string("model_path", "./model", "The path to export the model")
  6. flags.DEFINE_integer("export_version", 1, "Version number of the model")
  8. # Define the graph
  9. keys_placeholder = tf.placeholder(tf.int32, shape=[None, 1])
  10. keys = tf.identity(keys_placeholder)
  12. # Start the session
  14. # Export the model
  15. print("Exporting trained model to {}".format(FLAGS.model_path))
  16. model_exporter = exporter.Exporter(saver)
  17. model_exporter.init(
  18.   sess.graph.as_graph_def(),
  19.     named_graph_signatures={
  20.       'inputs': exporter.generic_signature({"keys": keys_placeholder, "features": inference_features}),
  21.       'outputs': exporter.generic_signature({"keys": keys, "softmax": inference_softmax, "prediction": inference_op})
  22.     })
  23. model_exporter.export(FLAGS.model_path, tf.constant(FLAGS.export_version), sess)
  24. print 'Done exporting!'

与SavedModel方法相比，两者都可以直接用TensorFlow Serving加载，我们使用deep_recommend_system导出两种模型方式测试过预测结果一模一样，只是模型文件大小不同。

导入带assert的模型文件

在NLP等场景除了参数文件，还需要导入vocabulary等文件，可以在exporter中设置assets_collection，参考 https://github.com/tensorflow/serving/issues/264 。

原文: http://docs.api.xiaomi.com/cloud-ml/modelservice/02_tensorflow_serving.html