parallel-processing - python : how to write this code to run on GPU?

Question

我已经尝试了很长一段时间来实现我的代码以在 GPU 上运行，但收效甚微。我真的很感激有人帮助实现。

让我就这个问题说几句。我有一个带有 N 个节点的图 G 和每个节点 x 上的分布 mx。我想计算所有边的每对节点的分布之间的距离。对于给定的对 (x,y)，我使用代码 ot.sinkhorn(mx, my, dNxNy)从 python POT 包计算距离。同样，mx, my 是节点 x 和 y 上大小为 Nx 和 Ny 的向量，dNxNy 是 Nx x Ny 距离矩阵。

现在，我发现此代码有一个 GPU 实现 ot.gpu.sinkhorn(mx, my, dNxNy) .然而，这还不够好，因为 I mx、my 和 dNxNy 需要在每次迭代时上传到 GPU，这是一个巨大的开销。因此，我们的想法是对 GPU 上的所有边进行并行处理。

代码的本质如下。 mx_all 是所有分布

for i,e in enumerate(G.edges):
    W[i] = W_comp(mx_all,dist,e)

def W_comp(mx_all, dist,  e):
    i = e[0]
    j = e[1]

    Nx = np.array(mx_all[i][1]).flatten()
    Ny = np.array(mx_all[j][1]).flatten()
    mx = np.array(mx_all[i][0]).flatten()
    my = np.array(mx_all[j][0]).flatten()

    dNxNy = dist[Nx,:][:,Ny].copy(order='C')

    W = ot.sinkhorn2(mx, my, dNxNy, 1)

下面是一个最小的工作示例。请忽略除虚线 === 之间的部分之外的所有内容迹象。

import ot
import numpy as np
import scipy as sc


def main():
    import networkx as nx

    #some example graph
    G = nx.planted_partition_graph(4, 20, 0.6, 0.3, seed=2)
    L = nx.normalized_laplacian_matrix(G)

    #this just computes all distributions (IGNORE)
    mx_all = []
    for i in G.nodes:
        mx_all.append(mx_comp(L,1,1,i))  

    #some random distance matrix (IGNORE)
    dist = np.random.randint(5,size=(nx.number_of_nodes(G),nx.number_of_nodes(G)))          

# ============================================================================= 
#this is what needs to be parallelised on GPU
    W = np.zeros(nx.Graph.size(G))
    for i,e in enumerate(G.edges):
        print(i)
        W[i] = W_comp(mx_all,dist,e)

    return W

def W_comp(mx_all, dist,  e):
    i = e[0]
    j = e[1]

    Nx = np.array(mx_all[i][1]).flatten()
    Ny = np.array(mx_all[j][1]).flatten()
    mx = np.array(mx_all[i][0]).flatten()
    my = np.array(mx_all[j][0]).flatten()

    dNxNy = dist[Nx,:][:,Ny].copy(order='C')

    return ot.sinkhorn2(mx, my, dNxNy,1)

# =============================================================================

#some other functions (IGNORE)
def delta(i, n):

    p0 = np.zeros(n)
    p0[i] = 1.

    return p0

# all neighbourhood densities
def mx_comp(L, t, cutoff, i):
    N = np.shape(L)[0]

    mx_all = sc.sparse.linalg.expm_multiply(-t*L, delta(i, N))
    Nx_all = np.argwhere(mx_all > (1-cutoff)*np.max(mx_all))

    return mx_all, Nx_all  

if __name__ == "__main__":
    main()  

谢谢!!

Answer 1

有一些软件包可以让您在 GPU 上运行代码。

您可以使用以下软件包之一:

pyCuda

numba(Pro)

Theano

当您想使用 numba 时，建议使用 Python Anaconda 发行版来执行此操作。此外，还需要 Anaconda Accelerate。您可以使用 conda install accelerate 安装它.在这个例子中，你可以看到 GPU 的使用是如何实现的 https://gist.githubusercontent.com/aweeraman/ae6e40f54a924f1f5832081be9521d92/raw/d6775c421aa4fa4c0d582e6c58873499d28b913a/gpu.py .
这是通过添加 target='cuda' 来完成的到 @vectorize装饰器。注意导入 from numba import vectorize . vectorize 装饰器将要加速的函数的签名作为输入。

祝你好运!

资料来源:

https://weeraman.com/put-that-gpu-to-good-use-with-python-e5a437168c01
https://www.researchgate.net/post/How_do_I_run_a_python_code_in_the_GPU