这篇文章主要讲解了如何使用keras和tensorflow保存为可部署的pb格式,内容清晰明了,对此有兴趣的小伙伴可以学习一下,相信大家阅读完之后会有帮助。

Keras保存为可部署的pb格式
加载已训练好的.h6格式的keras模型
传入如下定义好的export_savedmodel()方法内即可成功保存
import keras
import os
import tensorflow as tf
from tensorflow.python.util import compat
from keras import backend as K
def export_savedmodel(model):
'''
传入keras model会自动保存为pb格式
'''
model_path = "model/" # 模型保存的路径
model_version = 0 # 模型保存的版本
# 从网络的输入输出创建预测的签名
model_signature = tf.saved_model.signature_def_utils.predict_signature_def(
inputs={'input': model.input}, outputs={'output': model.output})
# 使用utf-8编码将 字节或Unicode 转换为字节
export_path = os.path.join(compat.as_bytes(model_path), compat.as_bytes(str(model_version))) # 将保存路径和版本号join
builder = tf.saved_model.builder.SavedModelBuilder(export_path) # 生成"savedmodel"协议缓冲区并保存变量和模型
builder.add_meta_graph_and_variables( # 将当前元图添加到savedmodel并保存变量
sess=K.get_session(), # 返回一个 session 默认返回tf的sess,否则返回keras的sess,两者都没有将创建一个全新的sess返回
tags=[tf.saved_model.tag_constants.SERVING], # 导出模型tag为SERVING(其他可选TRAINING,EVAL,GPU,TPU)
clear_devices=True, # 清除设备信息
signature_def_map={ # 签名定义映射
tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: # 默认服务签名定义密钥
model_signature # 网络的输入输出策创建预测的签名
})
builder.save() # 将"savedmodel"协议缓冲区写入磁盘.
print("save model pb success ...")
model = keras.models.load_model('model_data/weight.h6') # 加载已训练好的.h6格式的keras模型
export_savedmodel(model) # 将模型传入保存模型的方法内,模型保存成功.