如何将keras模型转换为协议缓冲区(.pb)文件?

jss*_*367 5 python protocol-buffers keras tensorflow

我有一个训练有素的keras模型,我想保存到协议缓冲区(.pb)文件中。当我这样做并加载模型时,预测是错误的(并且与原始模型不同)并且权重是错误的。这是模型类型:

type(model)
> keras.engine.training.Model
Run Code Online (Sandbox Code Playgroud)

这是我用来冻结并将其保存到.pb文件的代码。

from keras import backend as K
K.set_learning_phase(0)
import tensorflow as tf
from tensorflow.python.framework.graph_util import convert_variables_to_constants
keras_session = K.get_session()
graph = keras_session.graph
graph.as_default()
keep_var_names=None
output_names=[out.op.name for out in model.outputs]
clear_devices=True
with graph.as_default():
    freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference(keep_var_names or []))
    output_names = output_names or []
    output_names += [v.op.name for v in tf.global_variables()]
    input_graph_def = graph.as_graph_def()
    if clear_devices:
        for node in input_graph_def.node:
            node.device = ""
    frozen_graph = convert_variables_to_constants(keras_session, input_graph_def,
                                                  output_names, freeze_var_names)
tf.train.write_graph(frozen_graph, "model", "my_model.pb", as_text=False)
Run Code Online (Sandbox Code Playgroud)

然后我像这样阅读:

pb_file = 'my_model.pb'
with tf.gfile.GFile(pb_file, "rb") as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
with tf.Graph().as_default() as graph:
    tf.import_graph_def(graph_def)
ops = graph.get_operations()
def get_outputs(feed_dict, output_tensor):
    with tf.Session() as sess:

        sess.graph.as_default()
        tf.import_graph_def(graph_def, name='')

        output_tensor_loc = sess.graph.get_tensor_by_name(output_tensor)
        out = sess.run(output_tensor_loc, feed_dict=feed_dict)

        print("Shape is ", out.shape)

        return out
Run Code Online (Sandbox Code Playgroud)

然后,当我比较第一卷积层的权重时,它们的形状相同(形状看起来正确),但权重却不同。所有权重约为0:3,而在原始模型的同一层中,所有权重约为-256:256。

get_outputs(feed_dict, 'conv1_relu/Relu:0')
Run Code Online (Sandbox Code Playgroud)

上面的代码有问题吗?还是整个方法都不对?我在博客中看到有人在使用tf.train.Saver,但我没有这样做。我需要这样做吗?如果是这样,我该怎么做keras.engine.training.Model

rvi*_*nas 4

问:上面的代码有什么问题吗?或者这整个方法都是错误的?

答:主要问题是tf.train.write_graph保存了 TensorFlow 图,但没有保存模型的权重。


问:我需要使用吗tf.train.Saver?如果是这样,我该如何对我的模型执行此操作?

答:是的。除了保存图形(仅当您的后续脚本没有显式重新创建它时才需要)之外,您还应该用于tf.train.Saver保存模型的权重:

from keras import backend as K

# ... define your model in Keras and do some work

# Add ops to save and restore all the variables.
saver = tf.train.Saver()  # setting var_list=None saves all variables

# Get TensorFlow session
sess = K.get_session()

# save the model's variables
save_path = saver.save(sess, "/tmp/model.ckpt")
Run Code Online (Sandbox Code Playgroud)

调用saver.save还会保存 a MetaGraphDef,然后可以使用它来恢复图形,因此您没有必要使用tf.train.write_graph. 要恢复权重,只需使用saver.restore

with tf.Session() as sess:
  # restore variables from disk
  saver.restore(sess, "/tmp/model.ckpt")
Run Code Online (Sandbox Code Playgroud)

只要您使用 TensorFlow 后端(您仍然有 TensorFlow 图和权重),您使用 Keras 模型的事实就不会改变这种方法。有关在 TensorFlow 中保存和恢复模型的更多信息,请参阅保存和恢复教程。


保存 Keras 模型的替代(更简洁)方法

model.save('model_path.h5')现在,由于您使用的是 Keras 模型,因此按如下方式保存和恢复模型可能会更方便:

from keras.models import load_model

# restore previously saved model
model = load_model('model_path.h5')
Run Code Online (Sandbox Code Playgroud)

.pb更新:从.ckpt文件生成单个文件

如果您想生成单个.pb文件,请使用前一种tf.train.Saver方法。生成.ckpt文件(.meta保存图形和.data权重)后,您可以.pb通过调用Morgan 函数 freeze_graph来获取该文件,如下所示:

freeze_graph('/tmp', '<Comma separated output node names>') 
Run Code Online (Sandbox Code Playgroud)

参考:

  1. 在 TensorFlow 中保存和恢复
  2. StackOverflow 对TensorFlow 保存到文件/从文件加载图形的回答。
  3. 在 Keras 中保存/加载整个模型
  4. 摩根函数.pb从文件生成文件.ckpt