DGL v0.9 正式发布

自第一篇图卷积神经网络论文的发表已经过去了六年，图神经网络（GNN）已经成为重要的研究课题。研究人员正在积极探索更先进的GNN模型或训练方法。DGL开发团队密切关注新的研究趋势并基于此发布新特性。以下，我们列举了最新0.9版本的一些新功能。

使用cuGraph+DGL结合图分析算法和图神经网络

图神经网络（GNNs) 能够结合图数据的特征和结构信息，而经典的图分析算法则已被广泛用于刻画图结构，比如识别重要的节点和进行特征增强等。研究人员在探索更强大的图神经网络模型的过程中发现，当图神经网络和图分析技术组合时，它们的功能将被进一步扩展。为了达到这一目的，我们需要一个快速且可拓展的图分析引擎。Nvidia RAPIDS cuGraph library 是一款面向图分析算法的GPU算法库，支持多种经典的图分析算法，比如中心性计算（centrality) 和社区检测 (community detection) 等。“cuGraph能大大加速多种图分析算法，并在最新的NVIDIA GPUs上（RAPIDS支持Pascal 和更高版本的GPU架构）比在NetworkX 平均快1000倍。”

通过与NVidia工程师的合作，DGL 0.9版新增了两个API to_cugraph 和 from_cugraph，从而实现DGL图对象和cuGraph图对象之间的高效转换，使DGL用户可以访问cuGraph中的GPU图分析算法。

安装

我们推荐以下操作安装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 算法：基于Modularity Optimization对每个节点进行社区检测。
• Core number 算法：计算每个节点所属的k-core值。K-core子图中的每个顶点至少具有K的度数。一个节点的core值为K，则代表其最多属于一个K-core子图，而不属于任何一个K+1 - core子图。这两种算法刻画了一个节点的不同结构特征。Louvain算法将空间距离更近的节点归为一类，而Core Number算法则将局部结构相似的节点归为一类。下图展现了通过Louvain算法和Core Number算法分别在Zachary' s Karate Club Network 上的节点归类结果。

cuGraph为上述两种算法提供高效的GPU实现。我们可以使用 API‘to_cugraph’将‘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数据集准备节点特征。由于这两种算法都计算结构类别，我们将此类别转换为独热编码 (one-hot encoding) 并将结果作为为初始节点特征。

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 0.9版上线了对于PyTorch自动混合精度 (AMP) 的支持。PyTorch AMP用于混合精度训练，从而节省训练时间和GPU内存消耗。

用户可以通过torch.cuda.amp.autocast() 对模型的前向传递函数进行修饰，从而让PyTorch自动为每个op和tensor选择适当的数据类型。半精度训练可以有效减少内存占用，并允许系统使用 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
        

在使用半精度浮点型时，较小的梯度会被抹平并造成精度损失。PyTorch AMP提供的GradScaler模块通过放大loss值并对计算的梯度进行缩放来解决该问题。整套解决方案只需用户对代码进行如下修改：

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)

使用低精度或混合精度训练图神经网络仍然是当下活跃的研究课题。我们希望0.9版本能够促进该课题的更多研究。更多信息请查看文档Mixed Precision Training。

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

DGL-Go添加了对于整图预测任务的支持。它包括两个备受欢迎的GNN模型 - 图同构网络（GIN）和主邻域聚合（PNA），以及相关的训练和推理代码。例如，在ogbg-molpcba数据集上训练GIN模型，首先使用以下命令:

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

生成如下YAML配置文件。用户随后可以手动调整该配置文件。

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’命令来获取预先定义的模型结构。

dgl recipe get graphpred_pcba_gin.yaml

通过以下代码开始训练：

dgl train --cfg graphpred_ogbg-molpcba_gin.yaml

DGL-Go另一个新增功能是模型推理的支持。我们添加了‘dgl configure-apply’和‘dgl apply’两个命令用于在其他数据集上对已训练模型进行推理。例如，以下示例展示了使用在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文件保存，如下所示：

了解更多 如您想了解更多，请参阅本次版本的发布日志。