python - GPU 上的梯度计算

Question

在过去的两天里，我一直在尝试在 python 中实际计算图片堆栈的梯度。

Numpy.gradient() 给了我一些值，但它们对于我的边缘检测来说还不够好，所以我有点写了自己的内核来将数组中的每个值与灰度值相乘（我们称之为堆栈）。

所以 numpy.gradients() 非常快但不够好，我的梯度计算需要很多时间（大约 45 分钟）。在 OpenCV 和 CV Matrices 的帮助下在 C++ 中实现相同的功能使一切都变得更快（大约 3 分钟），但我无法在 python 和 c++ 之间构建一个包装器。

有没有办法在 GPU 上采用形状为 NxMxK 的 numpy float32 数组并以某种方式计算 x,y,z 方向的梯度？我试过 cudamat 但只有基本的矩阵运算，而且 tensorflow 或 pytorch 似乎很复杂:(

标签： pythonc++numpy