如何计算GPU上的成对距离矩阵

Question

我的代码中的瓶颈是我计算成对距离矩阵的区域 。 由于这是迄今为止最慢的部分，我花了很多时间来加速我的代码。

我发现许多加速在线使用文章，但收益微乎其微。 所以， 我正在寻找一种方法来使用我的GPU来创建距离矩阵 ，以进一步加快速度。 但是，我对使用GPU进行计算知之甚少。 任何人都可以帮我这样做吗？

在我的研究中，我发现了以下内容，但没有一个使用GPU ：

1. 这篇文章很有用，但加速很少。
2. 这篇文章提供了有关如何使用cython和numba的信息。

以下是如何计算成对距离矩阵的示例摘录：

import numpy as np
from scipy import spatial

rows = 1000
cols = 10
mat = np.random.randn(rows, cols)
d_mat = spatial.distance.cdist(mat, mat)

我的显卡是Nvidia Quadro M2000M

Answer 1

我能用这个：

import numpy as np
from numba import cuda

USE_64 = True

if USE_64:
    bits = 64
    np_type = np.float64
else:
    bits = 32
    np_type = np.float32

@cuda.jit("void(float{}[:, :], float{}[:, :])".format(bits, bits))
def distance_matrix(mat, out):
    m = mat.shape[0]
    n = mat.shape[1]
    i, j = cuda.grid(2)
    d = 0
    if i < m and j < m:
        for k in range(n):
            tmp = mat[i, k] - mat[j, k]
            d += tmp * tmp
        out[i, j] = d

def gpu_dist_matrix(mat):
    rows = mat.shape[0]

    block_dim = (16, 16)
    grid_dim = (int(rows/block_dim[0] + 1), int(rows/block_dim[1] + 1))

    stream = cuda.stream()
    mat2 = cuda.to_device(np.asarray(mat, dtype=np_type), stream=stream)
    out2 = cuda.device_array((rows, rows))
    distance_matrix[grid_dim, block_dim](mat2, out2)
    out = out2.copy_to_host(stream=stream)

    return out