简介
往期文章 我们给你推荐一种TensorFlow模型格式 介绍过， TensorFlow 官方推荐 SavedModel 格式作为在线服务的模型文件格式。近期 TensorFlow SavedModel 模块又推出了 simple_save 接口，简化了模型签名的构建和模型导出的成本，这期就结合 simple_tensorflow_serving 来做模型签名推荐以及快速上线相关的介绍。
新旧接口
回顾一下过去导出 SavedModel 的函数接口，由于一个模型可以有多 signature，每个签名可以有对应的 method name，因此我们需要引入 saved_model_builder、signature_constants、signature_def_utils、tag_constants 等变量，复制粘贴代码较多。
而新增的 simple_save()接口则简化很多，默认的签名就是 DEFAULT_SERVING _SIGNATURE_DEF_KEY，默认的 method 就是 PREDICT_METHOD_NAME，除了提供默认值用户也不需要调用 utils.build_tensor_info()来封装 op 了，简化代码如下。
可以看出，使用新的接口彻底解决了每次导出模型都要 Ctrl-c/Ctrl-v 的烦恼，接下来就介绍几个快速上线的模型。
Simplest 模型上线
Simplest 模型是我们用于 TensorFlow Session 性能测试的模型，最简化的模型签名就是一个 input op 和一个 output op，而且不混杂 add/minus/multiple/divide/convolution 等 kernel 的实现，理论上就可以测出 TensorFlow 本身与模型无关的性能，代码地址 tobegit3hub/tensorflow_examples 。
import tensorflow as tf
input_keys_placeholder = tf.placeholder(tf.int32, shape=[None, 1])output_keys = tf.identity(input_keys_placeholder)
session = tf.Session()tf.saved_model.simple_save(session, "./model/1", inputs={"keys": input_keys_placeholder}, outputs={"keys": output_keys})
复制代码
可以看出，简化后的代码只有 5 行，只需要 import tensorflow 即可，两个简单的 op，以及创建 Session 和导出模型。这个模型可以用 simple_tensorflow_serving 上线。
simple_tensorflow_serving --model_base_path="./model"
复制代码
模型上线后也可以预估，例如构造一个请求的 JSON，在浏览器、命令行或者任意编程语言实现的 HTTP 客户端请求即可。
Linear 模型上线
对于其他机器学习模型，模型导出方法也是类似的，首先是定义训练的 Graph，然后指定模型签名的 op，最后是调用 simple_save 接口来导出模型文件。这里以 Linear 模型为例，主要是训练数据可以在内存构造不需要依赖外网下载或者本地数据文件，完整代码在 tobegit3hub/tensorflow_examples 。
import numpy as npimport tensorflow as tf
# Prepare train datatrain_X = np.linspace(-1, 1, 100)train_Y = 2 * train_X + np.random.randn(*train_X.shape) * 0.33 + 10
# Define the modelX = tf.placeholder(tf.float32, shape=[2])Y = tf.placeholder(tf.float32, shape=[2])w = tf.Variable(0.0, name="weight")b = tf.Variable(0.0, name="bias")predict = X * w + b
loss = tf.square(Y - predict)train_op = tf.train.GradientDescentOptimizer(0.01).minimize(loss)
# Create session to runwith tf.Session() as sess: sess.run(tf.initialize_all_variables())
epoch = 1 for i in range(10): for (x, y) in zip(train_X, train_Y): _, w_value, b_value = sess.run([train_op, w, b], feed_dict={X: [x], Y: [y]}) print("Epoch: {}, w: {}, b: {}".format(epoch, w_value, b_value)) epoch += 1
export_dir = "./model/1/" print("Try to export the model in {}".format(export_dir)) tf.saved_model.simple_save(sess, export_dir, inputs={"x": X}, outputs={"y": predict})
复制代码
可以看出代码也不复杂，通过定义一个 X * W + b 的 op 来进行模型的预估，以及后续模型的预估，输入为 x，正好训练的 Graph 以及预估的 Graph 是重合的无序额外定义 op。这个模型上线方式是一样的，这里顺便介绍 simple_tensorflow_serving 的 code gen 的功能，我们可以在浏览器选择想要生成的编程语言就可以自动生成客户端代码。
simple_tensorflow_serving --model_base_path="./model"
复制代码
除此之外，在前端还有一个 JSON Inference 的功能，使用 code gen 生成的 JSON 请求示例，我们在前端就可以做深度学习模型的在线预估了。当然，这个请求数据是根据 TensorFlow SavedModel 的 Signature 来生成的，几乎所有的模型都可以这种方式自动生成请求数据以及客户端代码，Inference 结果或错误信息都会在浏览器中展示。
图像模型上线
前面已经介绍过通过 TensorFlow 模型的上线了，对于 SavedModel 模型来说，所有输入都是 Tensor，所谓"tensor in, tensor out"，这是通用 Serving + 任意 Model 的基础，但在绝大部分 CV 的场景下，模型的数据都是图片文件。
针对图像模型签名的优化，首先参考 Google 官方 TensorFlow Serving 的用法，例如 MNIST 模型的输入是[None, 2828]也是可以支持的，但要求客户端自己把 JPG、PNG 等图片文件转成 2828 的数组。然后我们参考 AWS 的 mxnet model server，用户不经可以通过 HTTP + application/json 的方式请求，还可以通过 form-data 来做预估，这样用户在浏览器上传的图片文件就可以直接请求到预估服务而不需要额外的转化。最终我们采用的是 base64 + decode op + reshape op 的方案，下面以 MNIST 模型为例。
import tensorflow as tffrom tensorflow.examples.tutorials.mnist import input_data
def inference(input): weights = tf.get_variable( "weights", [784, 10], initializer=tf.random_normal_initializer()) bias = tf.get_variable( "bias", [10], initializer=tf.random_normal_initializer()) logits = tf.matmul(input, weights) + bias
return logits
def main(): mnist = input_data.read_data_sets("./input_data")
x = tf.placeholder(tf.float32, [None, 784]) logits = inference(x) y_ = tf.placeholder(tf.int64, [None]) cross_entropy = tf.losses.sparse_softmax_cross_entropy( labels=y_, logits=logits) train_step = tf.train.GradientDescentOptimizer(0.5).minimize(cross_entropy)
init_op = tf.global_variables_initializer()
# Define op for model signature tf.get_variable_scope().reuse_variables()
model_base64_placeholder = tf.placeholder( shape=[None], dtype=tf.string, name="model_input_b64_images") model_base64_string = tf.decode_base64(model_base64_placeholder) model_base64_input = tf.map_fn(lambda x: tf.image.resize_images(tf.image.decode_jpeg(x, channels=1), [28, 28]), model_base64_string, dtype=tf.float32) model_base64_reshape_input = tf.reshape(model_base64_input, [-1, 28 * 28]) model_logits = inference(model_base64_reshape_input) model_predict_softmax = tf.nn.softmax(model_logits) model_predict = tf.argmax(model_predict_softmax, 1)
with tf.Session() as sess: sess.run(init_op)
for i in range(938): batch_xs, batch_ys = mnist.train.next_batch(64) sess.run(train_step, feed_dict={x: batch_xs, y_: batch_ys})
# Export image model export_dir = "./model/1" print("Try to export the model in {}".format(export_dir)) tf.saved_model.simple_save( sess, export_dir, inputs={"images": model_base64_placeholder}, outputs={ "predict": model_predict, "probability": model_predict_softmax })
复制代码
MNIST 的代码会稍微复杂一些，首先模型运行和模型签名用的 op 不同，因为我们希望预估的时候允许用户上传图片文件而不是训练时用的 Tensor，其次是要复用模型的权重用了 get_variable()以及 reuse_variables()，然后是在模型签名的输出 op 上做了各种 tf.decode_base64、tf.reshape、tf.argmax 等操作，这部分逻辑与训练无关却需要加到模型训练的脚本中方便导出模型。
通过上面的代码，我们就可以上线一个图像模型，可以接收以 form-data 形式传输的图片文件，图片文件的文件类型以及长、宽、通道数都没有要求，进入模型后会被统一处理成模型输入的 shape。而在 simple_tensorflow_serving 中，我们可以用到 Image Inference 功能，直接在浏览器页面上传图片文件，后台不需要额外处理直接请求 TensorFlow SavedModel 进行模型预测，结果返回到浏览器前端，这个步骤对于任意的 CV 场景以及 CV 模型都是通用的。
Estimator 模型上线
补充一下 TensorFlow Estimator 的模型导出和上线，Estimator 对 SavedModel 模型也进行了封装，用户使用 Estimator API 设置不需要指定模型签名，通过一个函数就可以导出 SavedModel，训练的输入就是模型预估时的输入，至于模型签名的 key 就是约定好的“inputs”，输出也是 Estimator 模型训练时的输出。
import tensorflow as tffrom tensorflow.contrib.learn.python.learn.utils import input_fn_utils
def input_fn(): features = {'a': tf.constant([["1"], ["2"]]), 'b': tf.constant([[3], [4]])} labels = tf.constant([0, 1]) return features, labels
feature_a = tf.contrib.layers.sparse_column_with_hash_bucket( "a", hash_bucket_size=1000)feature_b = tf.contrib.layers.real_valued_column("b")feature_columns = [feature_a, feature_b]
model = tf.contrib.learn.LinearClassifier(feature_columns=feature_columns)model.fit(input_fn=input_fn, steps=10)
feature_spec = tf.contrib.layers.create_feature_spec_for_parsing(feature_columns)serving_input_fn = input_fn_utils.build_parsing_serving_input_fn(feature_spec)savedmodel_path = "./model"model.export_savedmodel(savedmodel_path, serving_input_fn)
复制代码
如果我们去查看 SavedModel 的魔性签名，可以使用 tobegit3hub/tfmodel 或者 tobegit3hub/simple_tensorflow_serving ，可以看出输入只有一个 inputs，类型是 String，shape 是[None]。大家从 TensorFlow 官方的 wide and deep 代码中也可以看出，模型的输入是 adhoc 的字符串或者数字，在模型内部会做复杂的特征工程，因此输入不能使单一数值类型的，这里的 String 其实是要求用户根据每个特征的类型把样本序列化成 tf.train.Example，而这个 protobuf 序列化出来的 byte array 还不可被 ascii、unicode 编码成字符串，因此需要用 urlsafe_b64encode 转成可通过 HTTP 传输的字符串。注意，在 TensorFlow Serving 的 HTTP API 和 SimpleTensorFlowServing 的 HTTP API 中，都提供了额外的逻辑将 base64 编码的 tf.train.Example 转成 Tensor 再进行预估。
如果你正在考虑使用 TensorFlow Estimator 训练模型，导出模型和模型上线都是相当简单的，只是预估的客户端就比较复杂了，以前我写过的一个客户端数据生成程序仅供参考。
总结
最后总结下，本文只是介绍 SavedModel 模型签名的几种用法，并没有创造新的使用接口，但通过约定图像模型的输入字段以及 TensorFlow 内置 op 的预处理就可以实现更多“高级”的 Inference 接口。
大家在导出 TensorFlow SavedModel 时，根据预估接口在定义 inputs 和 outputs，即使输入 op 或输出 op 与训练的 Graph 不同也没关系，可以新增 op 并且通过 reuse_variables()等方式来共享权重实现更灵活的接口。而且一个 SavedModel 可以有多个 Signature，可以导出多个模型签名来保证接口的多样性，也可以使用本文介绍的 simple_save()接口简单地导出模型，或者使用 simple_tensorflow_serving 来快速上线和验证模型。
获授权转载自知乎账号：tobe
评论 1 条评论