c++ - 添加op tensorflow调试

Question

我想知道在向 tensorflow 添加新的 op 以进行调试时如何打印出输入张量的值。我一直在关注 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 数组，所以我可以打印出它的值。怎么了？

Answer 1

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