Tensorflow:创建tf.NodeDef()并设置属性

时间:2019-05-27 12:05:25

标签: python tensorflow protocol-buffers

我正在尝试创建新节点并设置其属性。

例如,打印一个图节点,我看到它的属性是:

attr {
  key: "T"
  value {
    type: DT_FLOAT
  }
}

我可以创建如下节点:

node = tf.NodeDef(name='MyConstTensor', op='Const',
                   attr={'value': tf.AttrValue(tensor=tensor_proto),
                         'dtype': tf.AttrValue(type=dt)})

但是如何添加key: "T"属性?即在这种情况下tf.AttrValue内应该是什么?

看着attr_value.proto,我已经尝试过:

node = tf.NodeDef()
node.name = 'MySub'
node.op = 'Sub'
node.input.extend(['MyConstTensor', 'conv2'])
node.attr["key"].s = 'T' # TypeError: 'T' has type str, but expected one of: bytes

更新:

我发现在Tensorflow中它应该写为:

node.attr["T"].type = b'float32'

但这会产生错误:TypeError: b'float32' has type bytes, but expected one of: int, long,我不确定哪个int值对应于float32。

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/attr_value.proto#L23

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/attr_value.proto#L35

2 个答案:

答案 0 :(得分:1)

通过反复试验,我发现它只是:

node.attr["T"].type = 1 # to set type to float32

答案 1 :(得分:0)

尝试将T作为字节传递:

node.attr["key"].s = b'T'

如果您想传递更多字符,请尝试字节数组类。

在AttrValue的protobuf定义中,将其定义为字节而不是字符串。 Protobuf manual指出这应该是python中的字符串,但是您的错误表明它更像是字节数组。