图神经网络（GNN）已成为处理各种机器学习图结构数据任务的重要工具。这些任务的范围从节点分类、链接预测到图形分类。它们还涵盖了广泛的应用，如社交网络分析、医疗保健中的药物发现、金融服务中的欺诈检测和分子化学。
在这篇文章中，我介绍了如何使用 cuGraph DGL，一个用于图计算的 GPU 加速库。它扩展了深度图库（DGL），这是一个支持大规模应用的 GNN 流行框架。
图神经网络基础
在深入研究 cuGraph DGL 之前，我想建立一些基础知识。GNN 是一种特殊的神经网络，设计用于处理结构化为图的数据。与传统的假设样本之间独立的神经网络不同，GNN 不能很好地与图形数据相匹配，它有效地利用了图形数据中丰富而复杂的互连。
简言之，GNN 通过多个步骤（通常称为层）在图结构中传播和转换节点特征来工作（图 1）。每个层基于其自身的特征和其邻居的特征来更新每个节点的特征。
在图 1 中，第一步“准备”一条由边缘及其连接节点的信息组成的消息，然后将消息“传递”给节点。该过程使模型能够学习节点、边和整个图的高级表示，可用于各种下游任务，如节点分类、链接预测和图分类。
图 2 显示了 2 层 GNN 应该如何计算节点 5 的输出。
处理大型图形时的瓶颈
GNN 采样和训练的瓶颈是缺乏能够扩展到处理数十亿甚至数万亿条边的现有实现，这种规模在现实世界的图问题中经常出现。例如，如果您正在处理具有数万亿条边的图形，则必须能够快速运行基于 DGL 的 GNN 工作流。
一种解决方案是使用 RAPIDS ，它已经拥有能够使用 GPU 扩展到数万亿边缘的基本元素。
什么是 RAPIDS cuGraph？
cuGraph 是 RAPIDS 人工智能生态系统的一部分，这是一套开源软件库，用于在 GPU 上执行端到端的数据科学和分析管道。cuGraph 库为图形分析提供了一个简单、灵活和强大的 API，使您能够以一定的规模和速度对图形数据进行计算。
什么是 DGL？
Deep Graph Library (DGL) 是一个 Python 库，其目标是通过提供直观的接口和高性能计算来简化图神经网络（GNN）的实现。
DGL 支持广泛的图形操作和结构，增强了复杂系统和关系的建模。它还与流行的深度学习框架（如 PyTorch 和 TensorFlow ）集成，促进了 GNN 的无缝开发和部署。
什么是 cuGraph DGL？
cuGraph DGL 是 cuGraph 的扩展，它与深度图库（DGL）集成，利用 GPU 的强大功能以前所未有的速度运行基于 DGL 的 GNN 工作流。这个库是 DGL 开发人员和 cuGraph 开发人员之间的协作。
除了 cuGraph DGL，cuGraph 还提供了 cuGraph 操作库，使 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 工作流包括以下步骤：
- 使用 cuGraph 操作模型，例如CuGraphSAGECon，代替了本地 DGL 模型(SAGEConv)。
- 创建CuGraphGraph对象。
- 使用cuGraph数据加载程序来代替本机 DGL 数据加载程序。
在 32 亿边缘图上使用 cugraph dgl，与单个 GPU UVA dgl 设置相比，当使用八个 GPU 进行采样和训练时，我们观察到速度提高了 3 倍。此外，当使用八个 GPU 进行采样和一个 GPU 进行训练时，我们看到了 2 倍的加速。
即将发布的一篇博客文章将提供更多关于收益和可扩展性的详细信息。
创建 cuGraph DGL 图形
创建cugraph_dgl图形直接从 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。
创建基于 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 的传统数据加载器。
import torchmetrics.functional as MF
import tempfile
import torch
def train(g, model):
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
features = g.ndata["feat"]["_N"].to("cuda")
labels = g.ndata["label"]["_N"].to("cuda")
train_nid = torch.tensor(range(g.num_nodes())).type(torch.int64)
temp_dir_name = tempfile.TemporaryDirectory().name
for epoch in range(10):
model.train()
sampler = cugraph_dgl.dataloading.NeighborSampler([10,10,10])
dataloader = cugraph_dgl.dataloading.DataLoader(g, train_nid, sampler,
batch_size=128,
shuffle=True,
drop_last=False,
num_workers=0,
sampling_output_dir=temp_dir_name)
total_loss = 0
for step, (input_nodes, seeds, blocks) in enumerate((dataloader)):
batch_inputs = features[input_nodes]
batch_labels = labels[seeds]
batch_pred = model(blocks, batch_inputs)
loss = F.cross_entropy(batch_pred, batch_labels)
total_loss += loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
sampler = cugraph_dgl.dataloading.NeighborSampler([-1,-1,-1])
dataloader = cugraph_dgl.dataloading.DataLoader(g, train_nid, sampler,
batch_size=1024,
shuffle=False,
drop_last=False,
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 中图神经网络简介，了解如何在 cuGraph PyG 生态系统中实现 GNN 的详细信息。
