NetworkX all_pairs_dijkstras 的 CuGraph 实现答案

【问题标题】：CuGraph implementation of NetworkX all_pairs_dijkstrasNetworkX all_pairs_dijkstras 的 CuGraph 实现
【发布时间】：2021-02-10 12:36:21
【问题描述】：

我正在尝试将我必须使用的 cpu 绑定算法转换为 GPU 算法，但我在使用 cugraph 时遇到了各种问题。其中一部分是我的无知，另一部分只是 cugraph 的起步阶段和不发达，最后一部分是我只是在寻找优雅的矢量化方法。

假设我有 m 个由 n 个特征组成的数据观察。我通过计算所有观测值的所有观测值的欧几里德距离来创建一个距离矩阵（注意：这不是我需要帮助的部分，也不是最佳的。只需添加这部分以获得易于理解、可重现的代码）

import numpy as np

def l1_distance(arr):
    return np.linalg.norm(arr, 1)

X = np.random.randint(low=0, high=255, size=(700,4096))

W = np.empty((700,700))

for i in range(700):
    for j in range(700):
        W[i,j] = l1_distance(X[i,:] - X[j,:])

第一次挑战

import networkx as nx

def build_weighted_graph(W):
    n = W.shape[0]
    Graph = nx.DiGraph()
        for i in range(n):
            for j in range(n):
                Graph.add_weighted_edges_from([(i,j,min(W[i,j], W[j,i]))])
    return Graph

其中输入 W 矩阵是欧几里得空间中的平方距离矩阵（节点最初由 x 特征组成）。例如。第 1 行 col 9 是节点 1 和节点 9 之间的距离，第 20 行 col 30 是节点 20 和节点 30 之间的距离，等等。该图现在绘制连接节点之间的边，边的权重是欧几里得距离测量.

我花了大约 8 个小时试图弄清楚如何将它转移到 GPU，但他们声称即使在 NVIDIA 自己的文档中

该代码块非常适合 NetworkX。然而，迭代数据帧并一次添加一个节点的过程对于 GPU 来说是有问题的，我们试图避免这种情况。 cuGraph 将数据存储在列（即数组）中。调整数组大小需要分配一个大一个元素的新数组，复制数据并添加新值。这不是很有效。

如果您的代码遵循上述在时间，我们建议重写该代码或按原样使用它在 NetworkX 中使用 cuGraph 加速算法。

所以我放弃了，让那部分保持原样。算法的下一部分使用dijkstra算法，计算所有节点到所有其他节点的最短路径

res = dict(nx.all_pairs_dijkstra_path_length(Graph))

在 cugraphs 实现中，它们只有单个源 dijkstra，它以图形和源节点作为参数。这与 networkx 的库形成对比，后者带有上述方法，并在所有节点上普遍应用 dijkstra。这意味着我必须为每个节点（更多的循环）迭代调用 SSSP（cugraph 的 dijkstra 的实现）。

在通过连接节点获得距离后，我创建另一个平方距离矩阵，该矩阵现在基于通过连接节点的距离，而不是最初采用欧几里德距离。

D = np.zeros([n,n])
    for i in range(n):
        for j in range(n):
            D[i,j] = res[i][j]

我这辈子都不知道如何为 GPU 向量化这些。任何朝着正确方向的帮助将不胜感激。目前对于我的算法运行的数据集，CPU 绑定算法需要大约 5 分钟才能运行 698 个节点。因此，为什么我要尝试在 GPU 方面加快速度

【问题讨论】：

请加minimal reproducible example，包括输入数据
在使用 GPU 之前，您是否尝试过从基于 python 的实现（例如 networkx）迁移到基于 C++ 的 python 库，例如 graph-tool？
我没有尝试过，因为我的工作的总体结果旨在扩大到 DGX GPU 集群并在海量数据集上运行（太大而无法在单台机器上运行）。所以这就是我努力让它准备好 GPU 的原因。
@SergeyBushmanov 添加了可重现的示例
您正在做的是 cuGraph 目前无法有效支持的事情，但解决这一不足是我们的路线图。 cuGraph 团队正在构建一个示例笔记本来回答您的问题，并将在此处发布。完成后我会添加一个带有链接的答案。

标签： gpu networkx dijkstra cupy rapids

【解决方案1】：

您的第一个问题（欧几里得距离矩阵的初始化）的答案似乎已在 Using cupy to create a distance matrix from another matrix on GPU 中得到解答，但您当然可以使用 cuDF 和 cuGraph 优化图形的创建和 Dijkstra 矩阵的计算。

为了有效地创建图形，您可以构建一个 cuDF 数据框，列出边及其权重。由于欧几里得距离矩阵的结构，这很简单。 cuGraph 将此数据框作为边列表接收并返回图形。然后，您可以遍历节点以计算到适用顶点的最短距离。如果问题规模增加，这可以在以后与 Dask 并行或分发。

对于这个问题大小，下面的代码比使用 nx.all_pairs_dijkstra_path_length 快大约 40 倍，它还包括初始距离计算。

import cupy as cp
import cudf as cd
import cugraph as cg

def build_weighted_graph_gpu(X, n):
    X_d = cp.array(X)
    G_d = cp.zeros([n, n])
    for i in range(n):
        G_d[i,:] = cp.abs(cp.broadcast_to(X_d[i,:], X_d.shape) - X_d).sum(axis=1)
    return G_d

def dijkstras_matrix_gpu(W_d):
    n = np.shape(W_d)[0]
    
    # Create a columnar dataframe describing edges
    df_g = cd.DataFrame({
        'source':      cp.array([x // n for x in range(n*n)]),
        'destination': cp.array([x % n for x in range(n*n)]),
        'weight':      cp.minimum(W_d, cp.transpose(W_d)).ravel()})
    
    graph_d = cg.Graph()
    
    graph_d.from_cudf_edgelist(df_g, source='source', destination='destination', edge_attr='weight', renumber=False)
    
    dist_d = cp.empty_like(W_d)

    for i in range(n):
        dist_d[i,:] = cp.asarray(cg.traversal.shortest_path(graph_d, i)['distance'])
    
    return dist_d

distance_matrix = dijkstras_matrix_gpu(build_weighted_graph_gpu(X))

【讨论】：