在Tensorflow2中将图形冻结为pb [英] Freezing graph to pb in Tensorflow2
问题描述
我们通过冻结图来保存它们,从而从TF1部署了许多模型:
We deploy lot of our models from TF1 by saving them through graph freezing:
tf.train.write_graph(self.session.graph_def, some_path)
# get graph definitions with weights
output_graph_def = tf.graph_util.convert_variables_to_constants(
self.session, # The session is used to retrieve the weights
self.session.graph.as_graph_def(), # The graph_def is used to retrieve the nodes
output_nodes, # The output node names are used to select the usefull nodes
)
# optimize graph
if optimize:
output_graph_def = optimize_for_inference_lib.optimize_for_inference(
output_graph_def, input_nodes, output_nodes, tf.float32.as_datatype_enum
)
with open(path, "wb") as f:
f.write(output_graph_def.SerializeToString())
然后通过以下方式加载它们:
and then loading them through:
with tf.Graph().as_default() as graph:
with graph.device("/" + args[name].processing_unit):
tf.import_graph_def(graph_def, name="")
for key, value in inputs.items():
self.input[key] = graph.get_tensor_by_name(value + ":0")
我们想以类似的方式保存TF2模型.一个protobuf文件,其中将包含图形和权重.我该如何实现?
We would like to save TF2 models in similar way. One protobuf file which will include graph and weights. How can I achieve this?
我知道有一些保存方法:
I know that there are some methods for saving:
-
keras.experimental.export_saved_model(model, 'path_to_saved_model')
这是实验性的,会创建多个文件:(.
Which is experimental and creates multiple files :(.
model.save('path_to_my_model.h5')
保存h5格式的:(.
tf.saved_model.save(self.model, "test_x_model")
再次保存多个文件:(.
Which agains save multiple files :(.
推荐答案
我使用TF2转换模型,例如:
I use TF2 to convert model like:
- 将
keras.callbacks.ModelCheckpoint(save_weights_only=True)
传递到model.fit
并在训练时保存checkpoint
; - 训练后,
self.model.load_weights(self.checkpoint_path)
加载checkpoint
,并转换为h5
:self.model.save(h5_path, overwrite=True, include_optimizer=False)
; - 将
h5
转换为pb
:
- pass
keras.callbacks.ModelCheckpoint(save_weights_only=True)
tomodel.fit
and savecheckpoint
while training; - After training,
self.model.load_weights(self.checkpoint_path)
loadcheckpoint
, and convert toh5
:self.model.save(h5_path, overwrite=True, include_optimizer=False)
; - convert
h5
topb
:
import logging
import tensorflow as tf
from tensorflow.compat.v1 import graph_util
from tensorflow.python.keras import backend as K
from tensorflow import keras
# necessary !!!
tf.compat.v1.disable_eager_execution()
h5_path = '/path/to/model.h5'
model = keras.models.load_model(h5_path)
model.summary()
# save pb
with K.get_session() as sess:
output_names = [out.op.name for out in model.outputs]
input_graph_def = sess.graph.as_graph_def()
for node in input_graph_def.node:
node.device = ""
graph = graph_util.remove_training_nodes(input_graph_def)
graph_frozen = graph_util.convert_variables_to_constants(sess, graph, output_names)
tf.io.write_graph(graph_frozen, '/path/to/pb/model.pb', as_text=False)
logging.info("save pb successfully!")
这篇关于在Tensorflow2中将图形冻结为pb的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持IT屋!