首页 > 其他 > 详细

【AI小疑问】h5转换pb模型(找不到输入到ReadVariableOp的变量)Cannot find the variable that is an input to the ReadVariableOp.

时间:2020-05-11 21:16:55      阅读:338      评论:0      收藏:0      [点我收藏+]

有人说要 更换Keras版本降级,感觉或许有用,但是为了统一版本又找到了一个解决方案,姑且不只还有什么坑

import keras.backend as K
k.set_learning_phase(0)

解决方案是将学习阶段设置为测试模式。

 

最终代码放一下

# -*- encoding: utf-8 -*-
"""
@File    : h5_to_pb.py
@Time    : 2020/5/11 17:54
@Author  : xuluda
@Email   : xuluda@163.com
@Software: PyCharm

将keras的.h5的模型文件,转换成TensorFlow的pb文件
"""
# ==========================================================

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


backend.set_learning_phase(0)


# from keras.models import Sequential

def h5_to_pb(h5_model, output_dir, model_name, out_prefix="output_", log_tensorboard=True):
    """.h5模型文件转换成pb模型文件
    Argument:
        h5_model: str
            .h5模型文件
        output_dir: str
            pb模型文件保存路径
        model_name: str
            pb模型文件名称
        out_prefix: str
            根据训练,需要修改
        log_tensorboard: bool
            是否生成日志文件
    Return:
        pb模型文件
    """
    if os.path.exists(output_dir) == False:
        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 = backend.get_session()

    from tensorflow.python.framework import graph_util, graph_io
    # 写入pb模型文件
    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 tensorflow.python.tools import import_pb_to_tensorboard
        import_pb_to_tensorboard.import_to_tensorboard(os.path.join(output_dir, model_name), output_dir)


if __name__ == ‘__main__‘:
    #  .h模型文件路径参数
    input_path = ‘/Users/jack/Documents/DaiCode/T4/‘
    weight_file = ‘false_positive3.h5‘
    weight_file_path = os.path.join(input_path, weight_file)
    output_graph_name = weight_file[:-3] + ‘.pb‘

    #  pb模型文件输出输出路径
    output_dir = osp.join(os.getcwd(), "trans_model")
    # model.save(xingren.h5)
    #  加载模型
    # h5_model = Sequential()
    h5_model = load_model(weight_file_path)
    # h5_model.save(weight_file_path)
    # h5_model.save(‘xingren.h5‘)
    h5_to_pb(h5_model, output_dir=output_dir, model_name=output_graph_name)
    print(‘Finished‘)

  

【AI小疑问】h5转换pb模型(找不到输入到ReadVariableOp的变量)Cannot find the variable that is an input to the ReadVariableOp.

原文:https://www.cnblogs.com/nowbug/p/12871148.html

(0)
(0)
   
举报
评论 一句话评论(0
关于我们 - 联系我们 - 留言反馈 - 联系我们:wmxa8@hotmail.com
© 2014 bubuko.com 版权所有
打开技术之扣,分享程序人生!