使用自定义操作冻结 TensorFlow 图

Question

我正在尝试使用标准构建图表。 TF 操作和自定义操作。单独执行它 python 运行良好，但另外我想导出这个图并重新加载它，但是 TF 在加载时难以解释我的自定义操作。

这是我导出图表的方式

export.py

import tensorflow as tf
from tensorflow.python.framework import graph_util
from tensorflow.python.framework import graph_io
with tf.device('/gpu:0'):
  with tf.Session() as sess:
    #build graph
    my_op = tf.load_op_library('/path/to/.so/')
    my_op = my_op.call_op
    math_op = tf.multiply(my_op(2),4)
    sess.run(math_op)

    #export graph
    oup_names = [None]
    oup_names[0] = sess.graph.get_operations()[-1].name
    constant_graph = graph_util.convert_variables_to_constants(sess,sess.graph.as_graph_def(),oup_names)
    graph_io.write_graph(constant_graph, "./","model.pb", as_text=False)

然后我尝试通过加载model.pb

import.py

import tensorflow as tf

with tf.gfile.GFile("./model.pb", "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, name="") #crashes here

#...

此时的错误信息：

在 import_graph_def raise ValueError('No op named %s in defined 操作。 % node.op) ValueError: No op named Example in defined 操作。

（顺便说一句。ExampleOp 是我的自定义操作类的名称）

如果我打印我收到的自定义操作的导出文本版本：

node {
  name: "Example/inp"
  op: "Const"
  device: "/device:GPU:0"
  attr {
    key: "dtype"
    value {
      type: DT_INT32
    }
  }
  attr {
    key: "value"
    value {
      tensor {
        dtype: DT_INT32
        tensor_shape {
        }
        int_val: 2
      }
    }
  }
}
node {
  name: "Example"
  op: "Example"
  input: "Example/inp"
  device: "/device:GPU:0"
}

我的猜测： TF 过度使用 op: "Example" 因为它缺乏关于如何使用此操作的定义（？）

对此有什么想法吗？

[更新]

我猜我的自定义操作的 bazel-BUILD 文件不完整，有没有一个示例如何在我的情况下编写一个？

Answer 1

通过添加解决它

my_op = tf.load_op_library('/path/to/.so/')

也在加载图表之前的 import.py 中