利用tf.saved_model導出keras模型

keras一般保存爲h5py格式的模型,當然也可以直接使用tf.saved_model保存爲pb模型,那如果想將保存的h5py模型導出爲pb模型該怎麼辦呢?以下代碼就可以完成該項功能。

假設我們保存了keras的模型爲model.json(結構)和weights.h5(權重),

首先讀取keras模型:

# tensorflow == 1.13.1
import tensorflow as tf

def load_keras_model(model_path, weights_path):
    fr = open(model_path, "r")
    model_json = fr.read()
    fr.close()
    model = tf.keras.models.model_from_json(model_json, custom_objects={"tf":tf})
    model.load_weights(weights_path)
    return model

然後轉換tensor name並導出模型:

model_export_dir = "./model/1"
model = load_keras_model("model.json", "weights.h5")
name_to_inputs = {i.name.split(":")[0]:i for i in model.inputs}
name_to_outputs = {i.name:i for i in model.outputs}
print(name_to_inputs)
print(name_to_outputs)
tf.saved_model.simple_save(tf.keras.backend.get_session(),
                           model_export_dir,
                           inputs=name_to_inputs,
                           outputs=name_to_outputs)

在model/1目錄下導出的模型結構爲:

saved_model.pb
variables |
          |variables.data-00000-of-00001
          |variables.index 

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章