注意
跳转到文末下载完整的示例代码。
使用 DGL 进行节点分类
GNN 是用于图上许多机器学习任务的强大工具。在本入门教程中，你将学习使用 GNN 进行节点分类的基本工作流程，即预测图中的节点类别。
完成本教程后，你将能够：
加载 DGL 提供的数据集。
使用 DGL 提供的神经网络模块构建 GNN 模型。
在 CPU 或 GPU 上训练和评估用于节点分类的 GNN 模型。
本教程假设你具有使用 PyTorch 构建神经网络的经验。
(预估时间：13 分钟)
import os
os.environ["DGLBACKEND"] = "pytorch"
import dgl
import dgl.data
import torch
import torch.nn as nn
import torch.nn.functional as F
GNN 节点分类概述
图数据上最流行且广泛采用的任务之一是节点分类，其中模型需要预测每个节点的真实类别。在图神经网络出现之前，许多提出的方法要么只使用连接性（如 DeepWalk 或 node2vec），要么简单地组合连接性和节点自身的特征。相比之下，GNN 提供了一个机会，可以通过组合连接性和 局部邻居 的特征来获得节点表示。
Kipf 等人的工作就是一个例子，它将节点分类问题表述为半监督节点分类任务。借助少量标记节点，图神经网络 (GNN) 可以准确预测其他节点的类别。
本教程将展示如何在 Cora 数据集上构建这样的 GNN 进行半监督节点分类，该数据集是一个引文网络，其中论文是节点，引文是边，并且只使用少量标记节点。任务是预测给定论文的类别。正如论文第 5.2 节所述，每篇论文节点都包含一个词频向量作为其特征，并经过归一化，使其总和为一。
加载 Cora 数据集
dataset = dgl.data.CoraGraphDataset()
print(f"Number of categories: {dataset.num_classes}")
Downloading /root/.dgl/cora_v2.zip from https://data.dgl.ai/dataset/cora_v2.zip...
/root/.dgl/cora_v2.zip: 0%| | 0.00/132k [00:00<?, ?B/s]
/root/.dgl/cora_v2.zip: 100%|██████████| 132k/132k [00:00<00:00, 7.74MB/s]
Extracting file to /root/.dgl/cora_v2_d697a464
Finished data loading and preprocessing.
NumNodes: 2708
NumEdges: 10556
NumFeats: 1433
NumClasses: 7
NumTrainingSamples: 140
NumValidationSamples: 500
NumTestSamples: 1000
Done saving data into cached files.
Number of categories: 7
一个 DGL Dataset 对象可以包含一个或多个图。本教程中使用的 Cora 数据集只包含一个图。
g = dataset[0]
DGL 图可以在两个类似字典的属性 ndata 和 edata 中存储节点特征和边特征。在 DGL Cora 数据集中，图包含以下节点特征：
train_mask: 一个布尔张量，指示节点是否在训练集中。val_mask: 一个布尔张量，指示节点是否在验证集中。test_mask: 一个布尔张量，指示节点是否在测试集中。label: 节点的真实类别。feat: 节点特征。
print("Node features")
print(g.ndata)
print("Edge features")
print(g.edata)
Node features
{'train_mask': tensor([ True, True, True, ..., False, False, False]), 'val_mask': tensor([False, False, False, ..., False, False, False]), 'test_mask': tensor([False, False, False, ..., True, True, True]), 'label': tensor([3, 4, 4, ..., 3, 3, 3]), 'feat': tensor([[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]])}
Edge features
{}
定义图卷积网络 (GCN)
本教程将构建一个两层图卷积网络 (GCN)。每一层通过聚合邻居信息来计算新的节点表示。
要构建多层 GCN，你可以简单地堆叠 dgl.nn.GraphConv 模块，这些模块继承自 torch.nn.Module。
from dgl.nn import GraphConv
class GCN(nn.Module):
def __init__(self, in_feats, h_feats, num_classes):
super(GCN, self).__init__()
self.conv1 = GraphConv(in_feats, h_feats)
self.conv2 = GraphConv(h_feats, num_classes)
def forward(self, g, in_feat):
h = self.conv1(g, in_feat)
h = F.relu(h)
h = self.conv2(g, h)
return h
# Create the model with given dimensions
model = GCN(g.ndata["feat"].shape[1], 16, dataset.num_classes)
DGL 提供了许多流行邻居聚合模块的实现。你只需一行代码即可轻松调用它们。
训练 GCN
训练这个 GCN 与训练其他 PyTorch 神经网络类似。
def train(g, model):
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
best_val_acc = 0
best_test_acc = 0
features = g.ndata["feat"]
labels = g.ndata["label"]
train_mask = g.ndata["train_mask"]
val_mask = g.ndata["val_mask"]
test_mask = g.ndata["test_mask"]
for e in range(100):
# Forward
logits = model(g, features)
# Compute prediction
pred = logits.argmax(1)
# Compute loss
# Note that you should only compute the losses of the nodes in the training set.
loss = F.cross_entropy(logits[train_mask], labels[train_mask])
# Compute accuracy on training/validation/test
train_acc = (pred[train_mask] == labels[train_mask]).float().mean()
val_acc = (pred[val_mask] == labels[val_mask]).float().mean()
test_acc = (pred[test_mask] == labels[test_mask]).float().mean()
# Save the best validation accuracy and the corresponding test accuracy.
if best_val_acc < val_acc:
best_val_acc = val_acc
best_test_acc = test_acc
# Backward
optimizer.zero_grad()
loss.backward()
optimizer.step()
if e % 5 == 0:
print(
f"In epoch {e}, loss: {loss:.3f}, val acc: {val_acc:.3f} (best {best_val_acc:.3f}), test acc: {test_acc:.3f} (best {best_test_acc:.3f})"
)
model = GCN(g.ndata["feat"].shape[1], 16, dataset.num_classes)
train(g, model)
In epoch 0, loss: 1.945, val acc: 0.122 (best 0.122), test acc: 0.160 (best 0.160)
In epoch 5, loss: 1.888, val acc: 0.560 (best 0.560), test acc: 0.545 (best 0.545)
In epoch 10, loss: 1.803, val acc: 0.618 (best 0.632), test acc: 0.623 (best 0.620)
In epoch 15, loss: 1.695, val acc: 0.652 (best 0.652), test acc: 0.639 (best 0.639)
In epoch 20, loss: 1.563, val acc: 0.678 (best 0.680), test acc: 0.676 (best 0.668)
In epoch 25, loss: 1.411, val acc: 0.708 (best 0.708), test acc: 0.700 (best 0.700)
In epoch 30, loss: 1.246, val acc: 0.714 (best 0.714), test acc: 0.713 (best 0.709)
In epoch 35, loss: 1.074, val acc: 0.728 (best 0.728), test acc: 0.742 (best 0.742)
In epoch 40, loss: 0.905, val acc: 0.740 (best 0.742), test acc: 0.752 (best 0.748)
In epoch 45, loss: 0.749, val acc: 0.750 (best 0.750), test acc: 0.759 (best 0.759)
In epoch 50, loss: 0.612, val acc: 0.752 (best 0.752), test acc: 0.766 (best 0.760)
In epoch 55, loss: 0.495, val acc: 0.748 (best 0.752), test acc: 0.768 (best 0.760)
In epoch 60, loss: 0.400, val acc: 0.760 (best 0.760), test acc: 0.767 (best 0.766)
In epoch 65, loss: 0.324, val acc: 0.762 (best 0.762), test acc: 0.771 (best 0.769)
In epoch 70, loss: 0.263, val acc: 0.768 (best 0.768), test acc: 0.776 (best 0.773)
In epoch 75, loss: 0.216, val acc: 0.766 (best 0.768), test acc: 0.776 (best 0.773)
In epoch 80, loss: 0.179, val acc: 0.764 (best 0.768), test acc: 0.771 (best 0.773)
In epoch 85, loss: 0.149, val acc: 0.764 (best 0.768), test acc: 0.769 (best 0.773)
In epoch 90, loss: 0.126, val acc: 0.766 (best 0.768), test acc: 0.769 (best 0.773)
In epoch 95, loss: 0.108, val acc: 0.764 (best 0.768), test acc: 0.768 (best 0.773)
在 GPU 上训练
在 GPU 上训练需要使用 to 方法将模型和图都放到 GPU 上，这与你在 PyTorch 中的操作类似。
g = g.to('cuda')
model = GCN(g.ndata['feat'].shape[1], 16, dataset.num_classes).to('cuda')
train(g, model)
下一步是什么？
# Thumbnail credits: Stanford CS224W Notes
# sphinx_gallery_thumbnail_path = '_static/blitz_1_introduction.png'
脚本总运行时间： (0 分 2.076 秒)