最簡單的方法—tensorflow 2.0 keras 保存 .pb 格式的模型

我們知道在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 即可。

然後,就沒有然後了...

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