深度图库 - Deep Graph Library (DGL) 文档

v0.9 版本亮点

在第一篇图卷积网络论文发表六年之后，研究人员仍在积极探索更高级的GNN架构或训练方法。作为DGL的开发团队，我们密切关注这些新的研究趋势，并发布新功能来促进它们的发展。本文重点介绍了v0.9新版本中的一些新功能。

使用 cuGraph+DGL 结合图分析与 GNN

图神经网络 (GNN) 能够结合图数据的特征和结构信息。当与图分析技术（例如特征增强）协同结合时，其能力可以进一步扩展。

图分析已被广泛用于表征图结构，例如识别重要节点，从而产生有趣的特征增强方法。为了利用这种协同作用，我们需要一个快速且可扩展的图分析引擎。NVIDIA 的 RAPIDS cuGraph 库提供了一系列用于图分析的 GPU 加速算法，例如中心性计算和社区检测。根据这份文档，“最新的 NVIDIA GPU（RAPIDS 支持 Pascal 及更高版本的 GPU 架构）使图分析比 NetworkX 平均快 1000 倍”。

与 NVIDIA 工程师合作，DGL v0.9 现在允许使用 to_cugraph 和 from_cugraph 这两个 API 在 DGLGraph 对象和 cuGraph 图对象之间进行转换，从而使 DGL 用户能够访问 cuGraph 中高效的图分析实现。

安装

要安装带有 PyTorch 和 DGL 的 cuGraph，我们建议遵循以下做法。Mamba 是 conda 的多线程版本。

conda install mamba -n base -c conda-forge

mamba create -n dgl_and_cugraph -c dglteam -c rapidsai-nightly -c nvidia -c pytorch -c conda-forge \
    cugraph pytorch torchvision torchaudio cudatoolkit=11.3 dgl-cuda11.3 tqdm

conda activate dgl_and_cugraph

通过 cuGraph 进行特征初始化

我们展示一个使用 cuGraph 提供的图分析算法进行节点特征初始化的示例。这里我们考虑两个选项

Louvain 算法，基于模块度优化检测每个节点的社区成员身份。
核心数算法，计算每个节点所属的最大 k-core 子图。图的 k-core 是包含度数大于或等于 k 的节点的最大子图。

这两种算法捕捉了节点不同的结构特征。Louvain 将空间距离接近的节点分组在一起，而核心数相同的节点在结构上彼此更相似。下图展示了在 Zachary's Karate Club Network 上由 Louvain 社区和核心数生成的节点着色。

node_coloring

cuGraph 提供了这两种算法的高效 GPU 实现。为了调用它们，我们使用 to_cugraph API 将 dgl.DGLGraph 转换为 cugraph.Graph。

import cugraph
import torch

def louvain(dgl_g):
    cugraph_g = dgl_g.to_cugraph().to_undirected()
    df, _ = cugraph.louvain(cugraph_g, resolution=3)
    # revert the node ID renumbering by cugraph
    df = cugraph_g.unrenumber(df, 'vertex').sort_values('vertex')
    return torch.utils.dlpack.from_dlpack(df['partition'].to_dlpack()).long()

def core_number(dgl_g):
    cugraph_g = dgl_g.to_cugraph().to_undirected()
    df = cugraph.core_number(cugraph_g)
    # revert the node ID renumbering by cugraph
    df = cugraph_g.unrenumber(df, 'vertex').sort_values('vertex')
    return torch.utils.dlpack.from_dlpack(df['core_number'].to_dlpack()).long()

通过 DGL 训练 GNN

然后我们使用上述函数为 ogbn-arxiv 数据集准备节点特征。请注意，由于这两种算法都计算结构类别，我们将它们转换为独热编码并连接起来作为初始节点特征。

import dgl.transforms as T
import torch.nn as nn
import torch.nn.functional as F

from dgl.nn import SAGEConv
from ogb.nodeproppred import DglNodePropPredDataset, Evaluator

device = torch.device('cuda')
dataset = DglNodePropPredDataset(name='ogbn-arxiv')
g, label = dataset[0]
transform = T.Compose([
    T.AddReverse(),
    T.AddSelfLoop(),
    T.ToSimple()
])
g = transform(g).int().to(device)
feat1 = louvain(g)
feat2 = core_number(g)
# convert to one-hot
feat1 = F.one_hot(feat1, feat1.max() + 1)
feat2 = F.one_hot(feat2, feat2.max() + 1)
# concat feat1 and feat2
x = torch.cat([feat1, feat2], dim=1).float()

然后我们训练一个简单的三层 GraphSAGE 模型（在此处查看完整的训练代码）。借助由图分析算法初始化的节点特征，我们能够仅使用纯结构信息在测试集上获得约 0.6 的准确率，这甚至优于使用原始输入节点特征的 MLP 模型。随着 DGL 新版本的发布，我们期待看到更多结合图分析的 GNN 创新。

FP16 与混合精度支持

DGL v0.9 现在完全兼容 PyTorch 自动混合精度 (AMP) 包进行混合精度训练，从而节省训练时间和 GPU 内存消耗。

通过使用 torch.cuda.amp.autocast() 包装前向传播，PyTorch 会自动为每个操作和张量选择适当的数据类型。半精度张量内存效率高，并且大多数基于半精度张量的操作更快，因为它们利用了 GPU 的 TensorCore。

import torch.nn.functional as F
from torch.cuda.amp import autocast

def forward(g, feat, label, mask, model):
    with autocast(enabled=True):
        logit = model(g, feat)
        loss = F.cross_entropy(logit[mask], label[mask])
        return loss

float16 格式的小梯度存在下溢问题（刷新为零）。PyTorch AMP 提供了一个 GradScaler 模块来解决这个问题。它将损失乘以一个因子，并在缩放后的损失上调用反向传播以防止下溢问题。然后在优化器更新参数之前对计算出的梯度进行反缩放。比例因子是自动确定的。

from torch.cuda.amp import GradScaler

scaler = GradScaler()

def backward(scaler, loss, optimizer):
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

总而言之，下面是示例。

import torch
import torch.nn as nn
from dgl.data import RedditDataset
from dgl.nn import GATConv
from dgl.transforms import AddSelfLoop

class GAT(nn.Module):
    def __init__(self, in_feats, num_classes, num_hidden=256, num_heads=2):
        super().__init__()
        self.conv1 = GATConv(in_feats, num_hidden, num_heads, activation=F.elu)
        self.conv2 = GATConv(num_hidden * num_heads, num_hidden, num_heads)

    def forward(self, g, h):
        h = self.conv1(g, h).flatten(1)
        h = self.conv2(g, h).mean(1)
        return h

device = torch.device('cuda')

transform = AddSelfLoop()
data = RedditDataset(transform)

g = data[0]
g = g.int().to(device)
train_mask = g.ndata['train_mask']
feat = g.ndata['feat']
label = g.ndata['label']
in_feats = feat.shape[1]

model = GAT(in_feats, data.num_classes).to(device)
model.train()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=5e-4)

for epoch in range(100):
    optimizer.zero_grad()
    loss = forward(g, feat, label, train_mask, model)
    backward(scaler, loss, optimizer)

使用低精度或混合精度训练 GNN 仍然是一个活跃的研究课题。我们希望新的 v0.9 版本将促进更多关于这个课题的研究。查看文档以了解更多信息。

DGL-Go 更新：模型推理和图预测

DGL-Go 现在支持训练 GNN 进行图属性预测任务。它包括两个流行的 GNN 模型——图同构网络 (GIN) 和主成分邻域聚合 (PNA)。例如，要在 ogbg-molpcba 数据集上训练 GIN 模型，首先使用命令生成 YAML 配置文件

dgl configure graphpred --data ogbg-molpcba --model gin

这将生成以下配置文件。用户可以手动调整配置文件。

version: 0.0.2
pipeline_name: graphpred
pipeline_mode: train
device: cpu                     # Torch device name, e.g., cpu or cuda or cuda:0
data:
    name: ogbg-molpcba
    split_ratio:                # Ratio to generate data split, for example set to [0.8, 0.1, 0.1] for 80% train/10% val/10% test. Leave blank to use builtin split in original dataset
model:
     name: gin
     embed_size: 300            # Embedding size
     num_layers: 5              # Number of layers
     dropout: 0.5               # Dropout rate
     virtual_node: false        # Whether to use virtual node
general_pipeline:
    num_runs: 1                 # Number of experiments to run
    train_batch_size: 32        # Graph batch size when training
    eval_batch_size: 32         # Graph batch size when evaluating
    num_workers: 4              # Number of workers for data loading
    optimizer:
        name: Adam
        lr: 0.001
        weight_decay: 0
    lr_scheduler:
        name: StepLR
        step_size: 100
        gamma: 1
    loss: BCEWithLogitsLoss
    metric: roc_auc_score
    num_epochs: 100             # Number of training epochs
    save_path: results          # Directory to save the experiment results

或者，用户可以获取原始实验中预定义超参数的模型配方。

dgl recipe get graphpred_pcba_gin.yaml

启动训练

dgl train --cfg graphpred_ogbg-molpcba_gin.yaml

另一个新增功能是用于对其他数据集上的已训练模型进行推理的新命令。例如，以下展示了如何将训练好的 ogbg-molpcba 上的 GIN 模型应用于 ogbg-molhiv

# Generate an inference configuration file from a saved experiment checkpoint
dgl configure-apply graphpred --data ogbg-molhiv --cpt results/run_0.pth

# Apply the trained model for inference
dgl apply --cfg apply_graphpred_ogbg-molhiv_pna.yaml

它将模型预测保存到 CSV 文件中，如下所示

csv_result

博客详情

v0.9 版本亮点

使用 cuGraph+DGL 结合图分析与 GNN

安装

通过 cuGraph 进行特征初始化

通过 DGL 训练 GNN

FP16 与混合精度支持

DGL-Go 更新：模型推理和图预测

更多阅读

关注我们

快速链接

资料

联系我们