使用dgl库实现GCN【官方案例】

学习目的

学习使用gnn进行节点分类的基本工作流程，即预测图中节点的类别。

关于GCN节点分类的综述

在图数据上最流行和广泛采用的任务之一是节点分类，其中模型需要预测每个节点的真实类别。
在图神经网络之前，许多被提出的方法要么单独使用连通性(如DeepWalk或node2vec)，要么简单地结合连通性和节点自身的特征。相比之下，gnn通过结合局部邻域的连通性和特征提供了获得节点表示的机会。Kipf等人将节点分类问题表述为半监督节点分类任务的例子。图神经网络(GNN)仅借助一小部分标记节点就能准确预测其他节点的类别。笔记将展示如何在Cora数据集上仅使用少量标签构建半监督节点分类的GNN，一个以论文为节点，引用为边的引文网络。任务是预测给定论文的类别。每个论文节点包含一个单词计数向量作为其特征，标准化后它们的总和为1。

导包

import os
os.environ["DGLBACKEND"] = "pytorch"
import dgl
import dgl.data
import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl.nn import GraphConv

os.environ[“DGLBACKEND”] = "pytorch"该行代码十分重要，因为它指定了DGL所运行的后端环境。

加载Cora数据集

dataset = dgl.data.CoraGraphDataset()
print(f"Number of categories: {dataset.num_classes}")

一个DGL数据集对象可以包含一个或多个图。本教程中使用的Cora数据集只包含一个图。

查看节点和边

"""
一个DGL数据集对象可以包含一个或多个图。本教程中使用的Cora数据集只包含一个图。
"""
g = dataset[0]

print("Node features")
print(g.ndata)
print("Edge features")
print(g.edata)

定义图卷积网络（GCN）

# 定义图卷积网络(GCN)
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

创建给定维度的模型

model = GCN(g.ndata["feat"].shape[1], 16, dataset.num_classes)

模型训练

使用DGL库训练网络类似于神经网络的训练，对于神经网络的训练步骤可以参考本博客：【转载】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 epoch in range(100):
        # 前向传播
        logits = model(g, features)

        # 计算预测值
        pred = logits.argmax(1)

        # 计算损失
        # 只计算训练集中节点的损失
        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})"
            )

在cuda上训练GCN

g = g.to('cuda') #在cuda上训练GCN
model = GCN(g.ndata['feat'].shape[1], 16, dataset.num_classes).to('cuda')
train(g, model)

省流版：完整代码

"""学习使用gnn进行节点分类的基本工作流程，即预测图中节点的类别。"""
import os
os.environ["DGLBACKEND"] = "pytorch"
import dgl
import dgl.data
import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl.nn import GraphConv
"""GNN节点分类综述
在图数据上最流行和广泛采用的任务之一是节点分类，其中模型需要预测每个节点的真实类别。
在图神经网络之前，许多被提出的方法要么单独使用连通性(如DeepWalk或node2vec)，
要么简单地结合连通性和节点自身的特征。

相比之下，gnn通过结合局部邻域的连通性和特征提供了获得节点表示的机会。
Kipf等人将节点分类问题表述为半监督节点分类任务的例子。图神经网络(GNN)
仅借助一小部分标记节点就能准确预测其他节点的类别。

本教程将展示如何在Cora数据集上仅使用少量标签构建半监督节点分类的GNN，一个以论文为节点，
引用为边的引文网络。任务是预测给定论文的类别。每个论文节点包含一个单词计数向量作为其特征，
标准化后它们的总和为1
"""


# 加载Cora数据集¶
dataset = dgl.data.CoraGraphDataset()
print(f"Number of categories: {dataset.num_classes}")

"""
一个DGL数据集对象可以包含一个或多个图。本教程中使用的Cora数据集只包含一个图。
"""
g = dataset[0]

print("Node features")
print(g.ndata)
print("Edge features")
print(g.edata)


# 定义图卷积网络(GCN)
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

# 创建具有给定维度的模型
model = GCN(g.ndata["feat"].shape[1], 16, dataset.num_classes)

# 模型训练
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 epoch in range(100):
        # 前向传播
        logits = model(g, features)

        # 计算预测值
        pred = logits.argmax(1)

        # 计算损失
        # 只计算训练集中节点的损失
        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})"
            )
g = g.to('cuda') #在cuda上训练GCN
model = GCN(g.ndata['feat'].shape[1], 16, dataset.num_classes).to('cuda')
train(g, model)