访问自定义操作的输入值时出现分段错误

Question

我只是关注the instructions，但在尝试读取我的 GPU 操作/上的输入值时，我总是会遇到段错误。如果我在 CPU 上执行相同的代码（然后使用不同的REGISTER_KERNEL_BUILDER），它会按预期工作。不幸的是，gdb 的回溯并没有给我更多信息，即使我使用 bazel 的调试标志构建了自定义操作。

这是我的代码

Interface.cc

REGISTER_OP("Interface")
    .Input("pointer_to_grid: int32")
    .Output("current_grid_data: float32")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
    shape_inference::ShapeHandle input_shape;
    TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &input_shape)); // allow only a 1D pointer address stored in an integer    
    return Status::OK();
    });

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

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

    printf("This works %d \n", input);
    printf("This does not %d \n", input(0)); //Segementation fault is here 

    //...

  }
};

REGISTER_KERNEL_BUILDER(Name("GridPointerInterface").Device(DEVICE_GPU), InterfaceGPU);

runme.py

import tensorflow as tf
import numpy as np
import sys
op_interface = tf.load_op_library('~/tensorflow/bazel-bin/tensorflow/core/user_ops/interface.so')
with tf.device("/gpu:0"):
  with tf.Session() as sess:
    sess.run(op_interface.interface_gpu(12))

我已经使用 TF 1.6 和 1.7 对其进行了测试。在我看来 TF 正在跳过内存分配，不幸的是我不知道如何强制这样做。

感谢您的建议

Answer 1

这是意料之中的，因为您正试图从 CPU 访问存储在 GPU 上的值（以便您可以打印它）。

在 GPU 上操作值的方法是通过 eigen。如果您查看 tensorflow 中其他内核的实现，您会看到诸如 output.flat<float32>().device(ctx->eigen_device<GPUDevice>()) = input.flat<float32>() + .... 之类的代码。这告诉 eigen 为你创建一个 cuda 内核。

如果您想直接操作 GPU 上的值，您需要同步 GPU 流并将其复制到 CPU 内存，这相当复杂。