图神经网络（GNN）已成为针对图结构数据的各种机器学习任务的强大工具。这些任务的范围从节点分类和链接预测到图分类。它们还涵盖广泛的应用，例如社交网络分析、医疗保健中的药物发现、金融服务中的欺诈检测以及分子化学。

在这篇文章中，我将介绍如何使用 cuGraph-DGL（一个用于图形计算的 GPU 加速库）。它扩展了深度图库 (DGL)，这是一种流行的 GNN 框架，可实现大规模应用。

图神经网络基础知识

在深入研究 cuGraph-DGL 之前，我想先了解一些基础知识。GNN 是一种特殊的神经网络，旨在处理图结构的数据。与假设样本之间独立的传统神经网络不同，GNN 不能很好地适应图数据，而 GNN 可以有效地利用图数据中丰富而复杂的互连。

简而言之，GNN 的工作原理是通过多个步骤（通常称为层）在图结构中传播和转换节点特征（图 1）。每层根据自己的特征和邻居的特征更新每个节点的特征。

在图 1 中，第一步“prepares准备”一条由来自边缘及其连接节点的信息组成的消息，然后将该消息“passes传递”到该节点。这个过程使模型能够学习节点、边和图作为一个整体的高级表示，这可用于各种下游任务，如节点分类、链接预测和图分类。

图 2 显示了 2 层 GNN 如何计算节点 5 的输出。

更新 2 层 GNN 中单个节点上的嵌入

处理大规模图时的瓶颈

GNN 采样和训练的瓶颈是缺乏可以扩展以处理数十亿甚至数万亿条边的现有实现，这种规模在现实世界的图问题中常见。例如，如果您正在处理具有数万亿条边的图，则必须能够快速运行基于 DGL 的 GNN 工作流程。

一种解决方案是使用 RAPIDS，它已经拥有能够使用 GPU 扩展到数万亿边缘的基本元素。

什么是 RAPIDS cuGraph？

cuGraph是 RAPIDS AI 生态系统的一部分，这是一个开源软件库套件，用于完全在 GPU 上执行端到端数据科学和分析管道。cuGraph 库为图分析提供了简单、灵活且强大的 API，使您能够大规模、快速地对图数据执行计算。

什么是DGL？

深度图库 (DGL)是一个 Python 库，旨在通过提供直观的界面和高性能计算来简化图神经网络 (GNN) 的实现。

DGL 支持广泛的图操作和结构，增强了复杂系统和关系的建模。它还与 PyTorch 和 TensorFlow 等流行的深度学习框架集成，促进 GNN 的无缝开发和部署。

什么是 cuGraph-DGL？

cuGraph-DGL 是 cuGraph 的扩展，它与深度图库 (DGL) 集成，利用 GPU 的强大功能以前所未有的速度运行基于 DGL 的 GNN 工作流程。该库是 DGL 开发人员和 cuGraph 开发人员之间的合作成果。

除了 cuGraph-DGL 之外，cuGraph 还提供 cugraph-ops 库，该库使 DGL 用户能够使用CuGraphSAGEConv、CuGraphGATConv和CuGraphRelGraphConv代替默认的SAGEConv、GATConv和RelGraphConv模型来获得性能提升。您还可以直接从库导入 SAGEConv、GATConv 和 RelGraphConv 模型cugraph_dgl。

在 GNN 采样和训练中，主要挑战是缺乏可以管理具有数十亿或数万亿条边的现实世界图问题的实现。为了解决这个问题，请使用 cuGraph-DGL，它具有使用 GPU 扩展到数万亿条边的固有能力。

设置 cuGraph-DGL

在深入研究代码之前，请确保您的 Python 环境中安装了 cuGraph 和 DGL。要安装启用 cuGraph-DGL 的环境，请运行以下命令：

conda install mamba -c conda-forge

mamba create -n cugraph_dgl_23_06 -c pytorch -c dglteam/label/cu118 -c rapidsai-nightly -c nvidia -c conda-forge dgl cugraph-dgl=23.10 pylibcugraphops=23.10 cudatoolkit=11.8 torchmetrics ogb

使用 cuGraph-DGL 实现 GNN

设置好环境后，将 cuGraph-DGL 付诸实践并构建一个简单的 GNN 以进行节点分类。将现有 DGL 工作流程转换为 cuGraph-DGL 工作流程有以下步骤：

1. 使用 cuGraph-ops 模型（例如 ）CuGraphSAGECon代替原生 DGL 模型 ( SAGEConv)。

2. 从 DGL 图创建CuGraphGraph对象。

3. 使用cuGraph数据加载器代替DGL 数据加载器。

在 32 亿边图上使用 cugraph-dgl，我们观察到，与单个 GPU UVA DGL 设置相比，使用八个 GPU 进行采样和训练时速度提高了 3 倍。此外，当使用 8 个 GPU 进行采样和 1 个 GPU 进行训练时，我们发现速度提高了 2 倍。

创建 cuGraph-DGL 图

要直接从 DGL 图创建cugraph_dgl图，请运行以下代码示例。

import dgl
import cugraph_dgl

dataset = dgl.data.CoraGraphDataset()
dgl_g = dataset[0]
# Add self loops as cugraph
# does not support isolated vertices yet
dgl_g = dgl.add_self_loop(dgl_g)
cugraph_g = cugraph_dgl.convert.cugraph_storage_from_heterograph(dgl_g, single_gpu=True)

有关创建cuGraph存储对象的更多信息，请参阅CuGraphStorage https://docs.rapids.ai/api/cugraph/stable/api_docs/api/cugraph-dgl/cugraph_dgl.cugraph_storage.cugraphstorage/#cugraph_dgl.cugraph_storage.CuGraphStorage。

创建基于 cuGraph-Ops 的模型

在此步骤中，唯一要做的修改是导入cugraph_ops基于 的模型。这些模型是上游模型例如dgl.nn.SAGECon.

# Drop in replacement for dgl.nn.SAGEConv
from dgl.nn import CuGraphSAGEConv as SAGEConv

import torch.nn as nn
import torch.nn.functional as F

class SAGE(nn.Module):
def __init__(self, in_size, hid_size, out_size):
super().__init__()
        self.layers = nn.ModuleList()
# three-layer GraphSAGE-mean
        self.layers.append(SAGEConv(in_size, hid_size, "mean"))
        self.layers.append(SAGEConv(hid_size, hid_size, "mean"))
        self.layers.append(SAGEConv(hid_size, out_size, "mean"))
        self.dropout = nn.Dropout(0.5)
        self.hid_size = hid_size
        self.out_size = out_size

def forward(self, blocks, x):
        h = x
for l_id, (layer, block) in enumerate(zip(self.layers, blocks)):
            h = layer(block, h)
if l_id != len(self.layers) - 1:
                h = F.relu(h)
                h = self.dropout(h)
return h


# Create the model with given dimensions
feat_size = cugraph_g.ndata["feat"]["_N"].shape[1]
model = SAGE(feat_size, 256, dataset.num_classes).to("cuda")

训练模型

在此步骤中，您选择使用cugraph_dgl.dataloading.NeighborSampler和cugraph_dgl.dataloading.DataLoader，替换上游 DGL 的传统数据加载器。

# Drop in replacement for dgl.nn.SAGEConv
from dgl.nn import CuGraphSAGEConv as SAGEConv

import torch.nn as nn
import torch.nn.functional as F

class SAGE(nn.Module):
def __init__(self, in_size, hid_size, out_size):
super().__init__()
        self.layers = nn.ModuleList()
# three-layer GraphSAGE-mean
        self.layers.append(SAGEConv(in_size, hid_size, "mean"))
        self.layers.append(SAGEConv(hid_size, hid_size, "mean"))
        self.layers.append(SAGEConv(hid_size, out_size, "mean"))
        self.dropout = nn.Dropout(0.5)
        self.hid_size = hid_size
        self.out_size = out_size

def forward(self, blocks, x):
        h = x
for l_id, (layer, block) in enumerate(zip(self.layers, blocks)):
            h = layer(block, h)
if l_id != len(self.layers) - 1:
                h = F.relu(h)
                h = self.dropout(h)
return h


# Create the model with given dimensions
feat_size = cugraph_g.ndata["feat"]["_N"].shape[1]
model = SAGE(feat_size, 256, dataset.num_classes).to("cuda")
                        num_workers=0,
                        sampling_output_dir=temp_dir_name)


      acc = evaluate(model, features, labels, dataloader)
print("Epoch {:05d} | Acc {:.4f} | Loss {:.4f} ".format(epoch, acc, total_loss))


def evaluate(model, features, labels, dataloader):
with torch.no_grad():
        model.eval()
        ys = []
        y_hats = []
for it, (in_nodes, out_nodes, blocks) in enumerate(dataloader):
with torch.no_grad():
                x = features[in_nodes]
                ys.append(labels[out_nodes])
                y_hats.append(model(blocks, x))
        num_classes = y_hats[0].shape[1]
return MF.accuracy(
            torch.cat(y_hats),
            torch.cat(ys),
            task="multiclass",
            num_classes=num_classes,
        )

train(cugraph_g, model)

Epoch 00000 | Acc 0.3401 | Loss 39.3890 
Epoch 00001 | Acc 0.7164 | Loss 27.8906 
Epoch 00002 | Acc 0.7888 | Loss 16.9441 
Epoch 00003 | Acc 0.8589 | Loss 12.5475 
Epoch 00004 | Acc 0.8863 | Loss 9.9894 
Epoch 00005 | Acc 0.8948 | Loss 9.0556 
Epoch 00006 | Acc 0.9029 | Loss 7.3637 
Epoch 00007 | Acc 0.9055 | Loss 7.2541 
Epoch 00008 | Acc 0.9132 | Loss 6.6912 
Epoch 00009 | Acc 0.9121 | Loss 7.0908

结论

通过将 GPU 加速的图计算的强大功能与 DGL 的灵活性相结合，cuGraph-DGL 成为处理图数据的任何人的宝贵工具。

这篇文章仅触及了 cuGraph-DGL 功能的冰山一角。我鼓励您进一步探索，尝试不同的 GNN 架构，并发现 cuGraph-DGL 如何加速基于图的机器学习任务。

有关如何在 cuGraph-PyG 生态系统中实现 GNN 的详细信息，请参阅我们的文章“ Intro to Graph Neural Networks with cuGraph-PyG https://medium.com/rapids-ai/intro-to-graph-neural-networks-with-cugraph-pyg-6fe32c93a2d0”。