趋近智
利用 PyTorch Geometric 组件(如 Data 对象和模块化 GNN 层),可以构建一个完整的图神经网络 (neural network)(GNN)训练流水线。我们编写了一个可重复使用的脚本,用于在节点分类任务上训练 GNN。该结构提供了一个标准模式,适用于许多基于图的项目。
我们将逐步讲解脚本的每个逻辑块:加载数据、定义模型、创建训练和评估循环,最后协调整个流程。
PyTorch Geometric 中的典型训练脚本遵循一个明确的结构,集成了数据处理、模型定义和训练过程。这种关注点分离使代码整洁、易读且易于修改。
标准 PyTorch Geometric 训练脚本的工作流,从数据加载到模型评估。
首先,我们导入必要的库。这包括用于核心张量操作的 torch,用于激活函数 (activation function)和损失函数 (loss function)的 torch.nn.functional,以及 torch_geometric 中的多个组件。我们将使用 Planetoid 数据集集合,其中包括 Cora、CiteSeer 和 PubMed 基准数据集。
import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv
# 加载 Cora 数据集
dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]
print(f'数据集: {dataset.name}')
print(f'图的数量: {len(dataset)}')
print(f'节点数量: {data.num_nodes}')
print(f'边数量: {data.num_edges}')
print(f'特征维度: {dataset.num_node_features}')
print(f'类别数量: {dataset.num_classes}')
Planetoid 对象会下载数据并将其处理为 PyG 可以使用的格式。由于 Cora 是单张图,数据集仅包含一个 Data 对象,我们可以通过 dataset[0] 访问它。这个 data 对象不仅包含节点特征 (data.x) 和图连接关系 (data.edge_index),还包含真实标签 (data.y) 以及用于半监督学习 (supervised learning) (semi-supervised learning)的预定义掩码 (data.train_mask、data.val_mask、data.test_mask)。
接下来,我们定义网络架构。我们创建一个继承自 torch.nn.Module 的类,这是在 PyTorch 中构建模型的标准方式。这使我们能够像使用其他 PyTorch 层一样使用 PyG 层。
在这里,我们将定义一个简单的双层 GCN。第一层 GCNConv 将输入节点特征映射到隐藏维度,第二层将隐藏表示映射到输出类别数。
class GCN(torch.nn.Module):
def __init__(self, num_features, num_classes):
super(GCN, self).__init__()
self.conv1 = GCNConv(num_features, 16)
self.conv2 = GCNConv(16, num_classes)
def forward(self, x, edge_index):
# 第一层 GCN
x = self.conv1(x, edge_index)
x = F.relu(x)
x = F.dropout(x, p=0.5, training=self.training)
# 第二层 GCN
x = self.conv2(x, edge_index)
# 对 NLLLoss 函数应用 LogSoftmax
return F.log_softmax(x, dim=1)
forward 方法定义了数据在网络中的流动方式。它接收节点特征矩阵 x 和 edge_index 作为输入。请注意,在第一层之后使用了 ReLU 激活函数 (activation function)和 Dropout 进行正则化 (regularization)。最终输出通过 log_softmax 函数处理,该函数通常与负对数似然损失 (F.nll_loss) 配对用于分类任务。
训练循环包含优化模型参数 (parameter)的逻辑。我们定义一个执行单次训练迭代的函数。
def train(model, optimizer, data):
model.train() # 将模型设置为训练模式
optimizer.zero_grad() # 清除梯度
# 执行单次前向传播
out = model(data.x, data.edge_index)
# 仅使用训练节点计算损失
loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
# 反向传播推导梯度
loss.backward()
# 更新参数
optimizer.step()
return loss.item()
该函数封装了标准的 PyTorch 训练步骤。对于 GNN 而言,最重要的是损失的计算方式。我们不是使用所有节点的输出,而是应用 data.train_mask 仅选择训练集中节点的 Logits 和标签。这是半监督节点分类的核心。
类似地,我们创建一个函数来评估模型在验证集或测试集上的表现。此函数不应更新模型权重 (weight),因此我们在 torch.no_grad() 下操作。
@torch.no_grad()
def test(model, data):
model.eval() # 将模型设置为评估模式
out = model(data.x, data.edge_index)
pred = out.argmax(dim=1) # 使用概率最高的类别
# 检查验证集和测试集
accs = []
for mask in [data.train_mask, data.val_mask, data.test_mask]:
correct = pred[mask] == data.y[mask] # 对比真实标签
accs.append(int(correct.sum()) / int(mask.sum())) # 计算准确率
return accs
在 test() 中,我们将模型设置为 eval() 模式,这会禁用 Dropout 等层。我们通过对输出 Logits 取 argmax 来获得所有节点的预测结果。然后,通过应用相应的掩码,计算每个数据拆分(训练、验证和测试)的准确率。
最后,我们编写主脚本逻辑来初始化所有组件并运行训练过程。我们实例化模型,定义 Adam 等优化器,然后进行指定轮数(epochs)的循环。在每一轮中,我们调用 train 和 test 函数并记录结果。
# 实例化模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GCN(dataset.num_features, dataset.num_classes).to(device)
data = data.to(device)
# 定义优化器
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
# 执行训练
best_val_acc = 0
for epoch in range(1, 201):
loss = train(model, optimizer, data)
train_acc, val_acc, test_acc = test(model, data)
# 简单的早停法
if val_acc > best_val_acc:
best_val_acc = val_acc
if epoch % 10 == 0:
print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, '
f'训练准确率: {train_acc:.4f}, 验证准确率: {val_acc:.4f}, 测试准确率: {test_acc:.4f}')
这个主代码块驱动了学习过程。通过跟踪验证准确率,您可以实现简单的早停法或保存验证集表现最好的模型。经过 200 轮训练后,在小型 Cora 数据集上训练的模型通常会在测试节点上获得较高的准确率。
GCN 在 Cora 数据集上的典型训练曲线。训练损失下降,而验证准确率上升并趋于稳定。
有了这个完整的脚本,您就拥有了一个强大且灵活的 GNN 训练模版。最后的动手操作环节将要求您应用此结构,从头开始在 Cora 数据集上训练和评估模型。
这部分内容有帮助吗?
© 2026 ApX Machine LearningAI伦理与透明度•