背景：目前keras框架使用簡單，很容易上手，深得廣大算法工程師的喜愛，但是當(dāng)部署到客戶端時，可能會出現(xiàn)各種各樣的bug，甚至不支持使用keras，本文來解決的是將keras的h5模型轉(zhuǎn)換為客戶端常用的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

#路徑參數(shù)

input_path = 'input path'

weight_file = 'weight.h5'

weight_file_path = osp.join(input_path,weight_file)

output_graph_name = weight_file[:-3] + '.pb'

#轉(zhuǎn)換函數(shù)

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')

將轉(zhuǎn)換成的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')

#預(yù)測結(jié)果

ret = sess.run(op, {input_x: 3, input_y: 4})

print(ret)

總結(jié)

以上是生活随笔為你收集整理的as转html5工具,将keras的h5模型转换为tensorflow的pb模型的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。