最简单的方法—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 即可。

然后,就没有然后了...

发布了9 篇原创文章 · 获赞 18 · 访问量 1万+

猜你喜欢

转载自blog.csdn.net/u012995500/article/details/103575879