Я довольно долго пытался реализовать свой код для работы на GPU, но без особого успеха.Буду очень признателен за помощь в реализации.
Позвольте мне сказать несколько слов о проблеме.У меня есть граф G с N узлами и распределением mx на каждом узле х.Я хотел бы вычислить расстояние между распределениями для каждой пары узлов для всех ребер.Для данной пары (x, y) я использую код ot.sinkhorn(mx, my, dNxNy) из пакета python POT для вычисления расстояния.Опять же, mx, my - это векторы размера Nx и Ny на узлах x и y, а 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()
Спасибо !!