访问自定义OP的输入值时的segmenation故障

Segmenation fault when accessing input value of custom op

本文关键字：segmenation 故障 输入 自定义 OP 访问      更新时间：2023-10-16

我只是按照说明进行操作，但是在尝试读取/gpu op上的输入值时，我总是最终会有一个segfault。如果我在CPU上执行相同的代码(然后使用不同的REGISTER_KERNEL_BUILDER(，则可以按预期工作。不幸的是，gdb的回溯没有给我进一步的信息，即使我使用Bazel的Debug Flag构建自定义OP。

这是我的代码

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＆amp;1.7。在我看来，TF正在跳过内存分配，不幸的是，我不确定如何强制这一点。

感谢您的建议

这是可以预期的，因为您正在尝试从CPU访问GPU上存储的值(因此您可以打印它(。

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

如果要直接操纵GPU上的值，则需要同步GPU流并将其复制到CPU内存，这很复杂。

相关文章：