使用自定义操作冻结TensorFlow图

Question

我正在尝试使用std构建图形。 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中引发ValueError（“在定义的操作中没有名为％s的操作。”％node.op）ValueError：在定义的操作中没有名为Example的操作。

（顺便说一句， 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"过于紧张，因为它缺少如何使用此操作的定义（？）

有什么想法吗？

[ 更新 ]

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

Answer 1

通过添加解决

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

在加载图形之前也在import.py中