添加op tensorflow调试

时间:2018-08-09 04:17:29

标签: c++ tensorflow

我想知道在为调试目的将新操作添加到tensorflow时如何打印出输入张量的值。我一直在使用cuda_op_kernel.cc来学习本教程,如下所示:

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include <typeinfo>

using namespace tensorflow;  // NOLINT(build/namespaces)

REGISTER_OP("AddOne")
    .Input("input: int32")
    .Output("output: int32")
    .Doc(R"doc(
Adds 1 to all elements of the tensor.

output: A Tensor.
  output = input + 1
)doc");

void AddOneKernelLauncher(const int* in, const int N, int* out);

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

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat<int32>();

    // Create an output tensor
    Tensor* output_tensor = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
                                                     &output_tensor));
    auto output = output_tensor->template flat<int32>();

    // Set all but the first element of the output tensor to 0.
    const int N = input.size();
    // Call the cuda kernel launcher
    std::cout << input.data() << std::endl;
    std::cout << input.data()[0] << std::endl;
    AddOneKernelLauncher(input.data(), N, output.data());
  }
};

REGISTER_KERNEL_BUILDER(Name("AddOne").Device(DEVICE_GPU), AddOneOp);

第二个std :: cout导致段错误。我的理解是input.data()应该是const int数组,所以我可以打印出它的值。怎么了?

1 个答案:

答案 0 :(得分:0)

我弄清楚了原因。 Input.data()是指向GPU中地址的指针。流执行器显然已经把它放在那里了。尝试在CPU上取消引用是不好的。