keras模型h5文件转pb文件（通过save_model保存的h5模型文件）

最新推荐文章于 2022-03-29 12:10:11 发布

ken_asr

最新推荐文章于 2022-03-29 12:10:11 发布

阅读量2.1k

点赞数

分类专栏：深度学习学习笔记文章标签： tensorflow 神经网络深度学习

本文链接：https://blog.csdn.net/zkgoup/article/details/105660089

版权

本文介绍如何将使用save_model方法保存的Keras .h5模型转换为TensorFlow的.pb文件，不适用于save_weights保存的模型。提供转换代码。

摘要由CSDN通过智能技术生成

该方法只适用于通过save_model保存的keras模型文件（.h5），不适用于通过save_weights保存的keras模型文件，直接上代码。

import keras
from keras.models import load_model
import tensorflow as tf
import os.path as osp
import os
from keras import backend as K


# 转换函数
def h5_to_pb(h5_model, output_dir, new_name, out_prefix="output_", log_tensorboard=True):
    if not osp.exists(output_dir):
        os.mkdir(output_dir)
    out_nodes = []
    for i in range(len(h5_model.outputs)):
        out_nodes.append(out_prefix + str(i + 1))
        tf.identity(h5_model.output[i], out_prefix + str(i + 1))
    sess = K.get_session()
    from tensorflow.python.framework import graph_util, graph_io
    init_graph = sess.graph.as_graph_def()
    main_graph = graph_util.convert_variables_to_constants(sess, init_graph, out_nodes)
    graph_io.write_graph(main_graph, output_dir, name=model_name, as_text=False)
    if log_tensorboard:
        from tenso