1

我对 CUDA/C++ 编程很陌生，我一直坚持将输入参数从 Tensorflow C++ API 传递给 CUDA 内核。

首先，我注册以下操作：

REGISTER_OP("Op")
.Attr("T: {float, int64}")
.Input("in: T")
.Input("angles: T")
.Output("out: T");

之后我想将第二个输入（角度）传递给 CPU/GPU 内核。不知何故，以下实现适用于 CPU 实现，但是当我在我的 GPU 上运行它时会在 Python 中引发错误... Python 错误消息：

Process finished with exit code -1073741819 (0xC0000005)

这就是我试图访问输入值的方式。请注意，“角度”的输入始终是单个值（float 或 int）：

void Compute(OpKernelContext* context) override {
...
const Tensor &input_angles = context->input(1);
auto angles_flat = input_angles.flat<float>();
const float N = angles_flat(0);
...
}

调用 CPU/GPU 内核如下：

...
Functor<Device, T>()(
            context->eigen_device<Device>(),
            static_cast<int>(input_tensor.NumElements()),
            input_tensor.flat<T>().data(),
            output_tensor->flat<T>().data(),
            N);
...

正如我之前所说，在 CPU 上运行这个 Op 就像我想要的那样工作，但是当我在 GPU 上运行它时，我总是得到上面提到的 Python 错误......有人知道如何解决这个问题吗？我只能猜测我正在尝试访问 GPU 上的错误地址angles_flat(0)......所以如果有人可以在这里帮助我，我将不胜感激！

4

0 回答 0