TensorFlow模型的簽名推薦與快速上線
簡介
往期文章 ofollow,noindex" target="_blank">我們給你推薦一種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 np
import tensorflow as tf
# Prepare train data
train_X = np.linspace(-1, 1, 100)
train_Y = 2 * train_X + np.random.randn(*train_X.shape) * 0.33 + 10
# Define the model
X = 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 run
with 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, 28*28]也是可以支援的，但要求客戶端自己把JPG、PNG等圖片檔案轉成28*28的陣列。然後我們參考AWS的mxnet model server，使用者不經可以通過HTTP + application/json的方式請求，還可以通過form-data來做預估，這樣使用者在瀏覽器上傳的圖片檔案就可以直接請求到預估服務而不需要額外的轉化。最終我們採用的是base64 + decode op + reshape op的方案，下面以MNIST模型為例。
import tensorflow as tf
from 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模型都是通用的。
總結
最後總結下，本文只是介紹SavedModel模型簽名的幾種用法，並沒有創造新的使用介面，但通過約定影象模型的輸入欄位以及TensorFlow內建op的預處理就可以實現更多“高階”的Inference介面。
大家在匯出TensorFlow SavedModel時，根據預估介面在定義inputs和outputs，即使輸入op或輸出op與訓練的Graph不同也沒關係，可以新增op並且通過reuse_variables()等方式來共享權重實現更靈活的介面。而且一個SavedModel可以有多個Signature，可以匯出多個模型簽名來保證介面的多樣性，也可以使用本文介紹的simple_save()介面簡單地匯出模型，或者使用simple_tensorflow_serving來快速上線和驗證模型。