如何在 GPU 上计算成对距离矩阵答案

【问题标题】：How to calculate pairwise distance matrix on the GPU如何在 GPU 上计算成对距离矩阵
【发布时间】：2017-10-09 22:26:30
【问题描述】：

我的代码中的瓶颈是我计算 pairwise distance matrix 的区域。由于这是迄今为止最慢的部分，我花了很多时间来加速我的代码。

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

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

This article 很有用，但加速效果微乎其微。
This article 提供了有关如何使用 cython 和 numba 的信息。

这是一个如何计算成对距离矩阵的示例 sn-p：

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

【问题讨论】：

我正在处理与您相同的问题 - 您是否有机会在过去几个月中取得任何突破，您可以在这里分享？
是的，我可以在 numba 库中使用 cuda 来编写代码，并且我得到了显着的加速...给我大约 24 小时，我可以找到我的代码并将其发布在这里一个答案。
@PaulTerwilliger 您还打算发布答案吗？在我努力为项目启动并运行一些持久的同源计算之前，我正在寻找同样的东西。
查看发布的答案

标签： python gpgpu distance-matrix

【解决方案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

【讨论】：

有兴趣查看共享内存的实现。干得好