别再死记硬背公式了！用Python+PyTorch Geometric（PyG）实战GCN，从消息传递到节点分类-深圳市維司達科技有限公司

用Python+PyTorch Geometric实战GCN：从零实现节点分类的工程指南

当第一次接触图卷积网络(GCN)时，很多人都会被复杂的数学公式和矩阵推导吓退。那些看似高深的理论背后，其实隐藏着一个简单的真相：GCN的核心思想不过是让节点学会从邻居那里收集信息。本文将带你用PyTorch Geometric(PyG)这个强大的图神经网络库，从零开始构建一个完整的GCN模型，并在Cora论文引用数据集上实现节点分类。

1. 环境准备与数据加载

在开始之前，确保你已经安装了Python 3.7+和以下库：

pip install torch torch-geometric torch-scatter torch-sparse torch-cluster torch-spline-conv -f https://data.pyg.org/whl/torch-1.10.0+cu113.html

PyTorch Geometric是专门为图神经网络设计的扩展库，它封装了常见的图操作，让我们可以专注于模型本身而不是底层实现。我们先加载Cora数据集——这是一个经典的论文引用网络：

from torch_geometric.datasets import Planetoid dataset = Planetoid(root='/tmp/Cora', name='Cora') data = dataset[0] print(f'Number of nodes: {data.num_nodes}') print(f'Number of edges: {data.num_edges}') print(f'Number of node features: {data.num_node_features}') print(f'Number of classes: {dataset.num_classes}')

你会看到输出显示Cora数据集包含2708个节点(论文)，5429条边(引用关系)，每个节点有1433维的特征(词袋表示)，共7个类别。

2. 理解GCN的消息传递机制

GCN的核心是消息传递框架，它包含三个关键步骤：

消息生成：每个节点为它的邻居准备要传递的消息
消息聚合：节点收集来自邻居的消息
节点更新：根据聚合后的消息更新节点表示

在PyG中，这可以通过定义一个MessagePassing类来实现。下面是一个简化版的GCN层实现：

import torch from torch_geometric.nn import MessagePassing from torch_geometric.utils import add_self_loops, degree class GCNConv(MessagePassing): def __init__(self, in_channels, out_channels): super().__init__(aggr='add') # 使用加法聚合 self.lin = torch.nn.Linear(in_channels, out_channels) def forward(self, x, edge_index): # 添加自环 edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0)) # 线性变换节点特征 x = self.lin(x) # 计算归一化系数 row, col = edge_index deg = degree(col, x.size(0), dtype=x.dtype) deg_inv_sqrt = deg.pow(-0.5) norm = deg_inv_sqrt[row] * deg_inv_sqrt[col] # 开始消息传递 return self.propagate(edge_index, x=x, norm=norm) def message(self, x_j, norm): return norm.view(-1, 1) * x_j

这个实现包含了GCN的所有关键要素：自环添加、度矩阵归一化和邻居信息聚合。message方法定义了如何准备消息，而聚合方式已经在初始化时设为'add'。

3. 构建完整的GCN模型

现在我们可以用上面定义的GCN层来构建一个完整的网络。一个典型的GCN架构包含2-3个图卷积层：

import torch.nn.functional as F from torch_geometric.nn import GCNConv class GCN(torch.nn.Module): def __init__(self, num_features, hidden_channels, num_classes): super().__init__() self.conv1 = GCNConv(num_features, hidden_channels) self.conv2 = GCNConv(hidden_channels, num_classes) def forward(self, x, 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)

这个模型的结构非常简单：

第一层GCN将1433维的特征映射到一个隐藏空间(通常设为16或64维)
ReLU激活函数引入非线性
Dropout防止过拟合
第二层GCN将隐藏表示映射到类别空间

4. 训练与评估模型

有了模型和数据，现在我们可以开始训练了。图神经网络的训练过程与传统神经网络类似：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = GCN(dataset.num_features, 16, dataset.num_classes).to(device) data = data.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4) def train(): 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() def test(): model.eval() out = model(data.x, data.edge_index) pred = out.argmax(dim=1) correct = pred[data.test_mask] == data.y[data.test_mask] acc = int(correct.sum()) / int(data.test_mask.sum()) return acc for epoch in range(1, 201): loss = train() if epoch % 10 == 0: acc = test() print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Acc: {acc:.4f}')

经过200轮训练，你应该能看到测试准确率稳定在80%左右。这已经是一个不错的结果，特别是考虑到我们只用了一个非常简单的模型。

5. 常见问题与调优技巧

在实际项目中，你可能会遇到以下问题及解决方案：

维度不匹配错误

RuntimeError: Sizes of tensors must match except in dimension 0

这通常是因为消息传递过程中特征维度不一致。检查：

所有GCN层的输入输出维度是否匹配
确保edge_index中的节点索引不超过实际节点数

过拟合问题

如果训练准确率很高但测试准确率低，可以尝试：

增加dropout率(0.5-0.8)
添加L2正则化(weight_decay)
减少GCN层数(通常2-3层足够)

梯度消失/爆炸

深层GCN容易出现梯度问题，解决方案：

使用残差连接
尝试不同的归一化方法(如BatchNorm)
降低学习率

一个更健壮的GCN实现可能长这样：

class BetterGCN(torch.nn.Module): def __init__(self, num_features, hidden_channels, num_classes): super().__init__() self.conv1 = GCNConv(num_features, hidden_channels) self.bn1 = torch.nn.BatchNorm1d(hidden_channels) self.conv2 = GCNConv(hidden_channels, num_classes) def forward(self, x, edge_index): x = self.conv1(x, edge_index) x = self.bn1(x) x = F.relu(x) x = F.dropout(x, p=0.6, training=self.training) x = self.conv2(x, edge_index) return F.log_softmax(x, dim=1)

6. 可视化节点嵌入

理解GCN如何学习的一个好方法是可视化节点的嵌入表示。我们可以使用t-SNE将高维嵌入降到2D：

from sklearn.manifold import TSNE import matplotlib.pyplot as plt def visualize(h, color): z = TSNE(n_components=2).fit_transform(h.detach().cpu().numpy()) plt.figure(figsize=(10,10)) plt.scatter(z[:, 0], z[:, 1], s=70, c=color, cmap="Set2") plt.show() model.eval() out = model(data.x, data.edge_index) visualize(out, color=data.y.cpu())

你会看到不同类别的节点在嵌入空间中形成了清晰的簇，这正是GCN的强大之处——它同时利用了节点特征和图结构信息。