將keras的h5模型轉換為tensorflow的pb模型
阿新 • • 發佈:2018-11-25
背景:目前keras框架使用簡單,很容易上手,深得廣大演算法工程師的喜愛,但是當部署到客戶端時,可能會出現各種各樣的bug,甚至不支援使用keras,本文來解決的是將keras的h5模型轉換為客戶端常用的tensorflow的pb模型並使用tensorflow載入pb模型。
h5_to_pb.py from keras.models import load_model import tensorflow as tf import os import os.path as osp from keras import backend as K #路徑引數 input_path = 'input path' weight_file = 'weight.h5' weight_file_path = osp.join(input_path,weight_file) output_graph_name = weight_file[:-3] + '.pb' #轉換函式 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.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 tensorflow.python.tools import import_pb_to_tensorboard import_pb_to_tensorboard.import_to_tensorboard(osp.join(output_dir,model_name),output_dir) #輸出路徑 output_dir = osp.join(os.getcwd(),"trans_model") #載入模型 h5_model = load_model(weight_file_path) h5_to_pb(h5_model,output_dir = output_dir,model_name = output_graph_name) print('model saved')
將轉換成的pb模型進行載入
load_pb.py import tensorflow as tf from tensorflow.python.platform import gfile def load_pb(pb_file_path): sess = tf.Session() with gfile.FastGFile(pb_file_path, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) sess.graph.as_default() tf.import_graph_def(graph_def, name='') print(sess.run('b:0')) #輸入 input_x = sess.graph.get_tensor_by_name('x:0') input_y = sess.graph.get_tensor_by_name('y:0') #輸出 op = sess.graph.get_tensor_by_name('op_to_store:0') #預測結果 ret = sess.run(op, {input_x: 3, input_y: 4}) print(ret)