Tensorflow:如何添加接受两个 1D vec 张量并输出标量的用户自定义操作？

Question

我正在尝试下面但没有工作。

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
using namespace tensorflow;
REGISTER_OP("Auc")
.Input("predicts: T1")
.Input("labels: T2")
.Output("z: double")
.Attr("T1: {float, double}")
.Attr("T2: {int32, int64}")
.SetIsCommutative()
.Doc(R"doc(
Given preidicts and labels output it's auc
)doc");

class AucOp : public OpKernel {
public:
explicit AucOp(OpKernelConstruction* context) : OpKernel(context) {}

void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& predicts_tensor = context->input(0);
    const Tensor& labels_tensor = context->input(1);
    auto predicts = predicts_tensor.flat<double>();
    auto labels = labels_tensor.flat<int32>();

    // Create an output tensor
    Tensor* output_tensor = NULL;
    TensorShape output_shape;
    OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output_tensor));

    output_tensor->flat<double>().setConstant(predicts(0) * labels(0));
}
};

REGISTER_KERNEL_BUILDER(Name("Auc").Device(DEVICE_CPU), AucOp);


test.py 

predicts = tf.constant([0.8, 0.5, 0.12])
labels = tf.constant([-1, 1, 1])

output = tf.user_ops.auc(predicts, labels)

with tf.Session() as sess:
    init = tf.initialize_all_variables()
    sess.run(init)

    print output.eval()

./测试.pyI tensorflow/core/common_runtime/local_device.cc:40] 本地设备内部运算并行线程数:8I tensorflow/core/common_runtime/direct_session.cc:60] Direct session inter op parallelism 线程:8F ./tensorflow/core/public/tensor.h:453] 检查失败:dtype() == DataTypeToEnum::v()(1 对 2)中止

Answer 1

问题是 predicts Python 程序中的张量类型为 float ，并且您的操作注册接受此作为 predicts 的有效类型输入(因为 T1 可以是 float 或 double )，但是 AucOp::Compute()假设 predicts输入始终具有类型 double (在调用 predicts_tensor.flat<double>() 时)。 tensorflow::Tensor当您请求不同类型的值时，类不会转换张量中元素的类型，而是会引发 fatal error 。

有几种可能的解决方案:

1. 要使事情快速进行，您可以更改 predicts 的类型在你的 Python 程序中到 tf.float64 (这是 Python 前端中 double 的同义词):

   predicts = tf.constant([0.8, 0.5, 0.12], dtype=tf.float64)
   

2. 您可以从定义一个仅接受单一类型输入的更简单的操作开始:

   REGISTER_OP("Auc")
   .Input("predicts: double")
   .Input("labels: int32")
   ...;
   

3. 您可以在 AucOp::Compute() 中添加代码方法来测试输入类型并根据需要访问输入值。 (使用 this->input_type(i) 查找第 i 个输入的类型。

4. 你可以定义一个模板类 AucOp<TPredict, TLabel> , 然后使用 TypeConstraint<>在REGISTER_KERNEL_BUILDER调用为预测和标签类型的四种有效组合中的每一种定义特化。这看起来像:

   REGISTER_KERNEL_BUILDER(Name("Auc")
                               .Device(DEVICE_CPU)
                               .TypeConstraint<float>("T1")
                               .TypeConstraint<int32>("T2"),
                          AucOp<float, int32>);
   // etc. for AucOp<double, int32>, AucOp<float, int64>, and AucOp<double, int64>.