如何合併兩個TensorFlow模型

雲水木石發表於2019-03-04

這是Tensorflow SavedModel模型系列文章的第三篇,也是終章。在《Tensorflow SavedModel模型的儲存與載入》中,我們談到了Tensorflow模型如何儲存為SavedModel格式,以及如何載入之。在《如何檢視tensorflow SavedModel格式模型的資訊》中,我們演示瞭如何檢視模型的signature和計算圖結構。在本文中,我們將探討如何合併兩個模型,簡單的說,就是將第一個模型的輸出,作為第二個模型的輸入,串聯起來形成一個新模型。

背景

為什麼需要合併兩個模型?

我們還是以《Tensorflow SavedModel模型的儲存與載入》中的程式碼為例,這個手寫數字識別模型接收的輸入是shape為[?, 784],這裡?代表可以批量接收輸入,可以先忽略,就把它固定為1吧。784是28 x 28進行展開的結果,也就是28 x 28灰度影像展開的結果。

問題是,我們送給模型的通常是圖片,可能來自檔案、可能來自攝像頭。讓問題變得複雜的是,如果我們通過HTTP來呼叫部署到伺服器端的模型,二進位制資料實際上是不方便HTTP傳輸的,這時我們通常需要對影像資料進行base64編碼。這樣伺服器端接收到的資料是一個base64字串,可模型接受的是二進位制向量。

很自然的,我們可以想到兩種解決方法:

  1. 重新訓練模型一個接收base64字串的模型。

    這種解決方法的問題在於:重新訓練模型很費時,甚至不可行。本文示例因為比較簡單,重新訓練也沒啥。如果是那種很深的卷積神經網路,訓練一次可能需要好幾天,重新訓練代價很大。更普遍的情況是,我們使用的是別人訓練好的模型,比如影像識別中普遍使用的Mobilenet、InceptionV3等等,都是Google、微軟這樣的公司,耗費大量的資源訓練出來的,我們沒有那個條件重新訓練。

  2. 在伺服器端增加base64到二進位制資料的轉換

    這種解決方法實現起來不復雜,但如果我們使用的是Tensorflow model server之類的方案部署的呢?當然我們也可以再開啟一個server,來接受客戶端的base64影像資料,處理完畢之後再轉發給Tensorflow model server,但這無疑增加了服務端的工作量,增加了服務端的複雜性。

在本文,我們將給出第三種方案:編寫一個Tensorflow模型,接收base64的影像資料,輸出二進位制向量,然後將第一個模型的輸出作為第二個模型的輸入,串接起來,儲存為一個新的模型,最後部署新的模型。

base64解碼Tensorflow模型

Tensorflow包含了大量影像處理和陣列處理的方法,所以實現這個模型比較簡單,模型包含了base64解碼、解碼PNG影像、縮放到28 * 28、最後展開為(1, 784)的陣列輸出,符合手寫數字識別模型的輸入,程式碼如下:

with tf.Graph().as_default() as g1:
  base64_str = tf.placeholder(tf.string, name=`input_string`)
  input_str = tf.decode_base64(base64_str)
  decoded_image = tf.image.decode_png(input_str, channels=1)
  # Convert from full range of uint8 to range [0,1] of float32.
  decoded_image_as_float = tf.image.convert_image_dtype(decoded_image,
                                                        tf.float32)
  decoded_image_4d = tf.expand_dims(decoded_image_as_float, 0)
  resize_shape = tf.stack([28, 28])
  resize_shape_as_int = tf.cast(resize_shape, dtype=tf.int32)
  resized_image = tf.image.resize_bilinear(decoded_image_4d,
                                           resize_shape_as_int)
  # 展開為1維陣列
  resized_image_1d = tf.reshape(resized_image, (-1, 28 * 28))
  print(resized_image_1d.shape)
  tf.identity(resized_image_1d, name="DecodeJPGOutput")

g1def = g1.as_graph_def()
複製程式碼

在該模型中,並不存在變數,都是一些固定的操作,所以無需進行訓練。

載入手寫識別模型

手寫識別模型參考《Tensorflow SavedModel模型的儲存與載入》一文,模型儲存在 “./model” 下,載入程式碼如下:

with tf.Graph().as_default() as g2:
  with tf.Session(graph=g2) as sess:
    input_graph_def = saved_model_utils.get_meta_graph_def(
        "./model", tag_constants.SERVING).graph_def

    tf.saved_model.loader.load(sess, ["serve"], "./model")

    g2def = graph_util.convert_variables_to_constants(
        sess,
        input_graph_def,
        ["myOutput"],
        variable_names_whitelist=None,
        variable_names_blacklist=None)
複製程式碼

這裡使用了g2定義了另外一個graph,和前面的模型的graph區分開來。注意這裡呼叫了graph_util.convert_variables_to_constants將模型中的變數轉化為常量,也就是所謂的凍結圖(freeze graph)操作。

在研究如何連線兩個模型時,我在這個問題上卡了很久。先的想法是合併模型之後,再載入變數值進來,但是嘗試之後,怎麼也不成功。後來的想法是遍歷手寫識別模型的變數,獲取其變數值,將變數值複製到合併的模型的變數,但這樣操作,使用模型時,總是提示有變數未初始化。

最後從Tensorflow模型到Tensorflow lite模型轉換中獲得了靈感,將模型中的變數固定下來,這樣就不存在變數的載入問題,也不會出現模型變數未初始化的問題。

執行convert_variables_to_constants後,可以看到有兩個變數轉化為了常量操作,也就是手寫數字識別模型中的w和b:

Converted 2 variables to const ops.
複製程式碼

連線兩個模型

利用tf.import_graph_def方法,我們可以匯入圖到現有圖中,注意第二個import_graph_def,其input是第一個graph_def的輸出,通過這樣的操作,就將兩個計算圖連線起來,最後儲存起來。程式碼如下:

with tf.Graph().as_default() as g_combined:
  with tf.Session(graph=g_combined) as sess:

    x = tf.placeholder(tf.string, name="base64_input")

    y, = tf.import_graph_def(g1def, input_map={"input_string:0": x}, return_elements=["DecodeJPGOutput:0"])

    z, = tf.import_graph_def(g2def, input_map={"myInput:0": y}, return_elements=["myOutput:0"])
    tf.identity(z, "myOutput")

    tf.saved_model.simple_save(sess,
              "./modelbase64",
              inputs={"base64_input": x},
              outputs={"myOutput": z})
複製程式碼

因為第一個模型不包含變數,第二個模型的變數轉化為了常量操作,所以最後儲存的模型檔案並不包含變數:

modelbase64/
├── saved_model.pb
└── variables

1 directory, 1 file

複製程式碼

測試

我們寫一段測試程式碼,測試一下合併之後模型是否管用,程式碼如下:

with tf.Session(graph=tf.Graph()) as sess:
  sess.run(tf.global_variables_initializer())

  tf.saved_model.loader.load(sess, ["serve"], "./modelbase64")
  graph = tf.get_default_graph()

  with open("./5.png", "rb") as image_file:
    encoded_string = str(base64.urlsafe_b64encode(image_file.read()), "utf-8")

  x = sess.graph.get_tensor_by_name(`base64_input:0`)
  y = sess.graph.get_tensor_by_name(`myOutput:0`)

  scores = sess.run(y,
           feed_dict={x: encoded_string})
  print("predict: %d, actual: %d" % (np.argmax(scores, 1), 5))
複製程式碼

這裡模型的輸入為base64_input,輸出仍然是myOutput,使用兩個圖片測試,均工作正常。

小結

最近三篇文章其實都是在研究我的微信小程式時總結的,為了更好的說明問題,我使用了一個非常簡單的模型來說明問題,但同樣適用於複雜的模型。

本文的完整程式碼請參考:github.com/mogoweb/aie…

希望這篇文章對您有幫助,感謝閱讀!同時敬請關注我的微信公眾號:雲水木石。

image

相關文章