将keras训练的h5模型转成pb格式

最新推荐文章于 2022-11-25 11:00:58 发布

dulingtingzi

最新推荐文章于 2022-11-25 11:00:58 发布

阅读量3.7k

点赞数 3

CC 4.0 BY-SA版权

分类专栏： deep learning tensorflow

本文链接：https://blog.youkuaiyun.com/dulingtingzi/article/details/90791869

deep learning 同时被 2 个专栏收录

44 篇文章

订阅专栏

tensorflow

7 篇文章

订阅专栏

本文介绍如何使用Keras和TensorFlow将训练好的深度学习模型转换为TF SavedModel格式，以便于部署和推理。文章详细展示了配置GPU、设置会话、加载模型权重以及导出模型的具体步骤。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

import os
import tensorflow as tf
from keras import backend as K
import logging
from tensorflow.python.util import compat
from keras.layers import Conv2D, Dense, Input, add, Activation, AveragePooling2D, GlobalAveragePooling2D
from keras.models import Model
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
K.set_session(tf.Session(config=config))
K.set_learning_phase(0)

def export_savedmodel(model,model_path='./'):
logging.info("Model input: {}, output: {}".format(model.input, model.output))
model_signature = tf.saved_model.signature_def_utils.predict_signature_def(
inputs={'input': model.input}, outputs={'output': model.output})

model_version = 1
export_path = os.path.join(
compat.as_bytes(model_path), compat.as_bytes(str(model_version)))
logging.info("Export the model to {}".format(export_path))

builder = tf.saved_model.builder.SavedModelBuilder(export_path)
builder.add_meta_graph_and_variables(
sess=K.get_session(),
tags=[tf.saved_model.tag_constants.SERVING],
clear_devices=True,
signature_def_map={
tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
model_signature
})
builder.save()
#下面是自定义的网络模型
img_input = Input(shape=input_shape)
output = densenet(img_input,num_classes)
model = Model(img_input, output)
model_path1='./saved_models/4_dzj_model.h5'
model.load_weights(model_path1)
model_path = ''
export_savedmodel(model, model_path)