使用自定义操作冻结TensorFlow图

时间:2018-04-04 14:44:24

标签: python tensorflow import export protocol-buffers

我正在尝试使用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的op)   操作“。 %node.op)ValueError:在定义中没有名为Example的op   操作

(顺便说一下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"
}

我猜: TFop: "Example"一起过度训练,因为它没有关于如何使用此操作的定义(?)

有关于此的任何想法吗?

[更新

我想我的自定义操作的bazel-BUILD文件是不完整的,是否有一个例子如何在我的情况下写一个?

1 个答案:

答案 0 :(得分:1)

通过添加

解决了这个问题
my_op = tf.load_op_library('/path/to/.so/')

import.py 之前加载图表