我們知道在tensorflow 2.0中利用 tensorflow.keras可以很方便地保存和加載 .h5 模型,但是 .h5格式的模型只適合在本地使用,不適合部署。tensorflow serving 提供了很強大的部署功能,但是僅支持 .pb 格式,關於 .pb 格式的好處請參考:TensorFlow 保存模型爲 PB 文件。本文說明如何用最簡單的方法利用 tensorflow.keras 保存爲 .pd 格式的模型。
當前的博客大都採用轉換的方式,轉換的過程中涉及到session,但實際上,session已經被新版的tensorflow摒棄了,所以就搞得很麻煩。原來的 的tf.keras.backend.get_session() 已經變成 tf.compat.v1.keras.backend.get_session(),可見這種方法有被拋棄的趨勢。
實際上,tensorflow.keras提供了一種非常簡單的保存和加載 .pb 格式模型的方法。廢話了這麼多,上代碼...
# 保存模型結構和參數到文件
tf.keras.models.save_model(network,"model_save_path") # 默認生成 .pb 格式模型,也可以通過save_format 設置 .h5 格式
print('模型已保存')
就是這麼簡單,保存之後會在 model_save_path 下生成如下三個文件。
加載模型也非常輕鬆愜意:
network=tf.keras.models.load_model("model_save_path")
要將 .h5 模型轉換成 .pb 格式同樣非常簡單,只需要先加載 .h5 再保存成 .pb 即可。
然後,就沒有然後了...