h5转pb的两个坑

1、需要加上如下设置,否则转换前后输出可能不一致,这个主要针对dropout、BN层训练测试不一致

      

from keras import backend as K
K.set_learning_phase(0) # 0 testing, 1 training mode

 

2、outputs而非output,否则会导致转换后无法 batch inference

     

def h5_to_pb(h5_model, output_dir, model_name, out_prefix="output_", log_tensorboard=True):
    if osp.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.outputs[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 tensorflow.python.tools import import_pb_to_tensorboard
        import_pb_to_tensorboard.import_to_tensorboard(osp.join(output_dir, model_name), output_dir)

 

posted @ 2019-12-11 14:29  牧马人夏峥  阅读(2241)  评论(0编辑  收藏  举报