是否有关于如何生成包含经过训练的 TensorFlow 图的 protobuf 文件的示例

Is there an example on how to generate protobuf files holding trained TensorFlow graphs

我正在研究 Google's example 如何在 Android 上部署和使用预训练的 Tensorflow 图(模型)。此示例使用 .pb 文件,位于:

https://storage.googleapis.com/download.tensorflow.org/models/inception5h.zip

这是一个 link 自动下载的文件

示例展示了如何将.pb文件加载到Tensorflow会话中并使用它来执行分类,但似乎没有提到如何生成这样的.pb文件,之后图已训练(例如,在 Python 中)。

有没有关于如何做到这一点的例子?

编辑: freeze_graph.py 脚本是 TensorFlow 存储库的一部分,现在用作生成代表 "frozen" 的协议缓冲区的工具经过训练的模型,来自现有的 TensorFlow GraphDef 和保存的检查点。它使用与下面描述的相同的步骤,但更容易使用。


目前这个过程没有很好的记录(并且有待完善),但大致步骤如下:

  1. 将您的模型构建并训练为 tf.Graph,称为 g_1
  2. 获取每个变量的最终值并将它们存储为 numpy 数组(使用 Session.run())。
  3. 在名为 g_2 的新 tf.Graph 中,使用相应 numpy 的值为每个变量创建 tf.constant() 张量在步骤 2 中获取的数组。
  4. 使用tf.import_graph_def() to copy nodes from g_1 into g_2, and use the input_map argument to replace each variable in g_1 with the corresponding tf.constant() tensors created in step 3. You may also want to use input_map to specify a new input tensor (e.g. replacing an input pipeline with a tf.placeholder())。使用 return_elements 参数指定预测输出张量的名称。

  5. 调用 g_2.as_graph_def() 获取图形的协议缓冲区表示。

(注意:生成的图在图中会有额外的节点用于训练。虽然它不是publicAPI的一部分,但你可能希望使用内部 graph_util.extract_sub_graph() 函数从图中去除这些节点。)

我不知道如何实现mrry描述的方法。但在这里我是如何解决它的。我不确定这是否是解决问题的最佳方法,但至少它解决了问题。

由于write_graph也可以存储常量的值,所以我在python中添加了以下代码,就在使用write_graph函数编写图形之前:

for v in tf.trainable_variables():
    vc = tf.constant(v.eval())
    tf.assign(v, vc, name="assign_variables")

这会创建常量,在训练后存储变量的值,然后创建张量“assign_variables”来分配它们到变量。现在,当您调用 write_graph 时,它会将变量的值以常量形式存储在文件中。

唯一剩下的部分是在 C 代码中调用这些张量“assign_variables”以确保您的变量被赋予存储在文件中的常量值。这是一种方法:

      Status status = NewSession(SessionOptions(), &session);
      std::vector<tensorflow::Tensor> outputs;
      char name[100];
      for(int i = 0;status.ok(); i++) {
        if (i==0)
            sprintf(name, "assign_variables");
        else
            sprintf(name, "assign_variables_%d", i);

        status = session->Run({}, {name}, {}, &outputs);
      }

刚找到这个 post 非常有用谢谢!我也使用@Mostafa 的方法,尽管我的 C++ 代码有点不同:

    std::vector<string> names;
    int node_count = graph.node_size();
    cout << node_count << " nodes in graph" << endl;

    // iterate all nodes
    for(int i=0; i<node_count; i++) {
        auto n = graph.node(i);
        cout << i << ":" << n.name() << endl;

        // if name contains "var_hack", add to vector
        if(n.name().find("var_hack") != std::string::npos) {
            names.push_back(n.name());
            cout << "......bang" << endl;
        }
    }
    session.Run({}, names, {}, &outputs);

注意我在 python

中使用 "var_hack" 作为我的变量名

这是对@Mostafa 回答的另一种看法。 运行 tf.assign 操作的更简洁的方法是将它们存储在 tf.group 中。这是我的 Python 代码:

  ops = []
  for v in tf.trainable_variables():
    vc = tf.constant(v.eval())
    ops.append(tf.assign(v, vc));
  tf.group(*ops, name="assign_trained_variables")

在 C++ 中:

  std::vector<tensorflow::Tensor> tmp;
  status = session.Run({}, {}, { "assign_trained_variables" }, &tmp);
  if (!status.ok()) {
    // Handle error
  }

这样,在 C++ 端,您只有一个 运行 的命名操作,因此您不必在遍历节点时搞得一团糟。

我在 Tensorflow 代码库中发现了一个 freeze_graph() 函数,在执行此操作时可能会有所帮助。据我了解,它在序列化 GraphDef 之前将变量与常量交换,因此当您从 C++ 加载此图时,它不再需要设置变量,您可以直接将其用于预测。

还有一个test for it and some description in the Guide.

这似乎是这里最简洁的选项。

替代我之前使用 freeze_graph() 的回答,这只有在您将其作为脚本调用时才有用,有一个非常好的函数可以为您完成所有繁重的工作并且适合调用来自您的正常模型训练代码。

convert_variables_to_constants() 做了两件事:

  • 它通过用常量替换变量来冻结权重
  • 它删除了与前馈预测无关的节点

假设 sess 是您的 tf.Session() 并且 "output" 是您的预测节点的名称,以下代码会将您的最小图形序列化为文本和二进制 protobuf。


from tensorflow.python.framework.graph_util import convert_variables_to_constants

minimal_graph = convert_variables_to_constants(sess, sess.graph_def, ["output"])

tf.train.write_graph(minimal_graph, '.', 'minimal_graph.proto', as_text=False)
tf.train.write_graph(minimal_graph, '.', 'minimal_graph.txt', as_text=True)