PyTorch Geometric 实现简易 Graph Convolutional Networks (GCN)

本文介绍如何使用 PyTorch Geometric 搭建一个简易的 Graph Convolutional Networks (GCN) 实例。我们将使用 PyTorch Geometric 库来定义一些图层，使用 Cora 数据集进行训练和测试。

导入相关库

import torch
import torch.nn.functional as F
import torch.nn as nn
import torch_geometric.nn as pyg_nn
from torch_geometric.datasets import Planetoid

加载 Cora 数据集

我们使用的是经典的 Cora 数据集，它是一个根据科学论文之间相互引用关系而构建的图数据集合，共有 2708 个节点和 10556 条边，每个节点的特征维度为 1433。

dataset = Planetoid(root='./data/Cora', name='Cora')
data = dataset[0]

定义 GraphConv 网络

GraphConv 层已经在 PyG 中预定义了，我们直接使用即可。这里我们使用两个 GraphConv 层，第一层的输入维度是每个节点的特征维度，输出维度为 16，第二个层的输入维度为 16，输出维度为节点类别数。每个层后面跟一个 ReLU 和 dropout 操作。

class GraphConv(nn.Module):
    def __init__(self, num_node_features, num_classes):
        super(GraphConv, self).__init__()
        self.conv1 = pyg_nn.GraphConv(num_node_features, 16)
        self.conv2 = pyg_nn.GraphConv(16, num_classes)
        
    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)
        
        return F.log_softmax(x, dim=1)

定义模型、损失函数和优化器

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 设备
epochs = 200 # 学习轮数
lr = 0.0003 # 学习率
num_node_features = dataset.num_node_features # 每个节点的特征数
num_classes = dataset.num_classes # 每个节点的类别数

model = GraphConv(num_node_features, num_classes).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
loss_function = nn.NLLLoss()

模型训练

model.train() # 将模型设为训练模式
for epoch in range(epochs):
    optimizer.zero_grad() # 梯度清零
    pred = model(data) # 前向传播
    
    # 计算损失和精度
    loss = loss_function(pred[data.train_mask], data.y[data.train_mask]) 
    correct_count_train = pred.argmax(axis=1)[data.train_mask].eq(data.y[data.train_mask]).sum().item()
    acc_train = correct_count_train / data.train_mask.sum().item()
    
    loss.backward() # 反向传播
    optimizer.step() # 更新参数
    
    # 输出训练过程信息
    if epoch % 10 == 0:
        print(f"Epoch: {epoch}, Training Loss: {loss.item()}, Training Accuracy: {acc_train}")

模型训练完成后，我们需要进行模型验证，即使用测试集来测试模型的泛化能力。在Cora数据集中，我们使用掩码机制来划分训练集和测试集。

模型验证

model.eval()

pred = model(data)

训练集（使用了掩码）

correct_count_train = pred.argmax(axis=1)[data.train_mask].eq(data.y[data.train_mask]).sum().item()
acc_train = correct_count_train / data.train_mask.sum().item()
loss_train = loss_function(pred[data.train_mask], data.y[data.train_mask]).item()

测试集

correct_count_test = pred.argmax(axis=1)[data.test_mask].eq(data.y[data.test_mask]).sum().item()
acc_test = correct_count_test / data.test_mask.sum().item()
loss_test = loss_function(pred[data.test_mask], data.y[data.test_mask]).item()

print('Train Accuracy: {:.4f}'.format(acc_train), 'Train Loss: {:.4f}'.format(loss_train))
print('Test Accuracy: {:.4f}'.format(acc_test), 'Test Loss: {:.4f}'.format(loss_test))

最终输出的结果如下所示：

Train Accuracy: 1.0000 Train Loss: 0.3787
Test Accuracy: 0.3790 Test Loss: 1.7529

可以看出，模型在训练集上的准确率为1，而在测试集上的准确率只有0.379，存在较大的过拟合现象，这需要我们进行更多的优化和改进。

测试集

本文通过使用Pytorch Geometric来实现了一个简单的Graph Convolutional Network（GCN）模型，以Cora数据集为例，实现了节点分类任务。在实现的过程中，我们介绍了Pytorch Geometric的基本使用方法，并详细讲解了GCN的实现原理。

在实际应用中，GCN作为一种图神经网络，在社交网络、推荐系统等领域具有广泛的应用前景，希望读者通过本文的介绍，可以更好地了解GCN，并在实践中深入掌握其使用方法。