Keras模型.h5转成tensorflow的.pb

在keras中，保存.h5模型可能出现两种保存方法：

即（1）保存模型权重和网络结构，以及（2）只保存模型权重。

.h5转.pb的两种方式：

方法一：

Keras的.h5模型转成tensorflow的.pb格式模型，方便后期的前端部署。直接上代码

1
 

1. from keras.models import Model

2. from keras.layers import Dense, Dropout

3. from keras.applications.mobilenet import MobileNet

4. from keras.applications.mobilenet import preprocess_input

5. from keras.preprocessing.image import load_img, img_to_array

6. import tensorflow as tf

7. from keras import backend as K

8. import os

9.  
10.  

11. base_model = MobileNet((None, None, 3), alpha=1, include_top=False, pooling='avg', weights=None)

12. x = Dropout(0.75)(base_model.output)

13. x = Dense(10, activation='softmax')(x)

14.  

15. model = Model(base_model.input, x)

16. model.load_weights('mobilenet_weights.h5')

17.  

18. def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True):

19. from tensorflow.python.framework.graph_util import convert_variables_to_constants

20. graph = session.graph

21. with graph.as_default():

22. freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference(keep_var_names or []))

23. output_names = output_names or []

24. output_names += [v.op.name for v in tf.global_variables()]

25. input_graph_def = graph.as_graph_def()

26. if clear_devices:

27. for node in input_graph_def.node:

28. node.device = ""

29. frozen_graph = convert_variables_to_constants(session, input_graph_def,

30. output_names, freeze_var_names)

31. return frozen_graph

32.  

33. output_graph_name = 'NIMA.pb'

34. output_fld = ''

35. #K.set_learning_phase(0)

36.  

37. print('input is :', model.input.name)

38. print ('output is:', model.output.name)

39.  

40. sess = K.get_session()

41. frozen_graph = freeze_session(K.get_session(), output_names=[model.output.op.name])

42.  

43. from tensorflow.python.framework import graph_io

44. graph_io.write_graph(frozen_graph, output_fld, output_graph_name, as_text=False)

45. print('saved the constant graph (ready for inference) at: ', os.path.join(output_fld, output_graph_name)

方法二：

先建立好网络结构model，然后调用load_weights函数加载权重参数，然后再开始转：

本文主要记录Keras训练得到的.h5模型文件转换成TensorFlow的.pb文件

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
#*-coding:utf-8-*
"""
将keras的.h5的模型文件，转换成TensorFlow的pb文件
"""
# ==========================================================
from keras.models import load_model
import tensorflow as tf
import os
from keras import backend
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 = 'satellite/train_dir/models/'
weight_file = 'satellite_iv3_ft.h5'
weight_file_path = os.path.join(input_path, weight_file)
output_graph_name = weight_file[:-3] + '.pb'
#
pb模型文件输出输出路径
output_dir = os.path.join(os.getcwd(), "satellite/train_dir/models/")
#
加载模型
h5_model = load_model(weight_file_path)
h5_to_pb(h5_model, output_dir=output_dir, model_name=output_graph_name)
print('Finished')

其中，load_model()这个函数不是从keras中加载的，功能就是自己创建的网络结构，创建好网络结构后加载模型，最后才得到h5_model。  然后才开始调用h5_to_pb转化。

最后

以上就是迅速嚓茶最近收集整理的关于Keras模型.h5转成tensorflow的.pb的全部内容，更多相关Keras模型.h5转成tensorflow内容请搜索靠谱客的其他文章。