Python：如何编写此代码以在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)。但是，这还不够好，因为我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 上运行代码。

您可以使用以下包之一：

1. pyCuda
2. numba（专业版）
3. 西阿诺

当您想使用 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。矢量化装饰器将要加速的函数的签名作为输入。

祝你好运！

资料来源：

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