python - TensorFlow 接受哪些 Python 类型作为 "tensor"类型的 Attr？

Question

我正在 C++ 中定义一个新 Op，它接受 tensor 类型的单个属性，大致如下 these instructions 。操作码的精简版本如下:

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"

using namespace tensorflow;

REGISTER_OP("DoStuff")
    .Attr("attr: tensor = { dtype: DT_FLOAT }")
    .Input("in: float")
    .Output("out: float");

class DoStuffOp : public OpKernel {
public:
    explicit DoStuffOp(OpKernelConstruction *context) : OpKernel(context) {
        OP_REQUIRES_OK(context, context->GetAttr("attr", &attr_));
        // ...
    }

    void Compute(OpKernelContext *context) override {
        // ...
    }

private:
    Tensor attr_;
};

REGISTER_KERNEL_BUILDER(Name("DoStuff").Device(DEVICE_CPU), DoStuffOp);

我可以将 Op 编译成 .so 文件。但是，我不知道如何成功传递 attr 的值。当我在 Python 中运行以下命令时:

import tensorflow as tf
dostufflib = tf.load_op_library('build/do_stuff.so')
sess = tf.InteractiveSession()

A = [[1.0, 2.0, 3.0],
     [4.0, 5.0, 6.0]]
X = tf.Variable(tf.constant(1.0))

Y = dostufflib.do_stuff(X, A)

我收到TypeError:不知道如何将参数“attr”的[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]转换为TensorProto。我所做的一切似乎都无法满足类型转换:list、numpy array、tf.Tensor、tf.Variable等等。如何将 Python 变量作为张量属性传递到 Op 中？

Answer 1

经过更多的寻找，我发现了tf.contrib.util.make_tensor_proto ，一个将 python 标量、python 列表、numpy ndarray 或 numpy 标量转换为 tf.TensorProto 的函数目的。作品如下:

A = tf.contrib.util.make_tensor_proto([[1.0, 2.0, 3.0],[4.0, 5.0, 6.0]])
X = tf.Variable(tf.constant(1.0))

Y = dostufflib.do_stuff(X, A)