圖神經網絡（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 dglimport cugraph_dgldataset = 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.SAGEConvfrom dgl.nn import CuGraphSAGEConv as SAGEConvimport torch.nn as nnimport torch.nn.functional as Fclass 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 dimensionsfeat_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 MFimport tempfileimport torchdef 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 的詳細信息。
?
