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)

 


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM