我只是跟着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 正在跳过内存分配,不幸的是我不确定如何强制执行此操作。
谢谢你的建议
最佳答案
这是预料之中的,因为您正试图从 CPU 访问存储在 GPU 上的值(因此您可以打印它)。
在 GPU 上操作值的方法是通过 eigen。如果你查看 tensorflow 中其他内核的实现,你会看到类似 output.flat<float32>().device(ctx->eigen_device<GPUDevice>()) = input.flat<float32>() + ....
的代码。 .这会告诉 eigen 为您创建一个 cuda 内核。
如果您想直接操作 GPU 上的值,您需要同步 GPU 流并将它们复制到 CPU 内存,这相当复杂。
关于python - 访问自定义操作的输入值时出现段错误,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/49558207/