訪問自定義操作的輸入值時出現段錯誤

Question

我只是按照說明進行操作，但是在嘗試讀取我的GPU操作上的輸入值時，我總是會遇到段錯誤。 如果我在CPU上執行相同的代碼（然后使用另一個REGISTER_KERNEL_BUILDER ），它將按預期工作。 不幸的是，即使我使用bazel的debug標志構建了自定義op， gdb的回溯也無法提供更多信息。

這是我的代碼

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上操縱值的方法是通過特征。 如果您在tensorflow中查看其他內核的實現，則會看到諸如output.flat<float32>().device(ctx->eigen_device<GPUDevice>()) = input.flat<float32>() + .... 這告訴本征為您創建一個cuda內核。

如果要直接操縱GPU上的值，則需要同步GPU流並將其復制到CPU內存，這相當復雜。