PyTorch Geometric入门教程：从零开始掌握图神经网络

2026-04-08 09:19:54作者：袁立春Spencer

概念解析

核心概念速览

如何判断你的数据是否适合图神经网络处理？现实世界中许多数据具有天然的图结构特性，比如社交网络中的用户关系、分子结构中的原子连接等。图神经网络（GNN）正是处理这类数据的强大工具，它能够有效捕捉节点之间的依赖关系和全局结构信息。

PyTorch Geometric（PyG）是基于PyTorch的图神经网络库，专为简化图深度学习任务而设计。它提供了丰富的图数据处理工具和GNN模型实现，让开发者能够快速构建和训练图神经网络模型。

图数据结构详解

图数据由哪些基本组件构成？在PyG中，图数据主要通过Data对象来表示，包含以下核心组件：

节点特征(x)：形状为[num_nodes, num_features]的张量，存储每个节点的特征信息
边索引(edge_index)：形状为[2, num_edges]的COO格式张量，定义节点之间的连接关系
边特征(edge_attr)：可选的边属性张量，存储与边相关的特征
节点标签(y)：节点的类别标签，用于节点分类任务

下面是创建一个简单图数据对象的示例：

from torch_geometric.data import Data
import torch

# 节点特征：3个节点，每个节点2个特征
x = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=torch.float)
# 边索引：4条边，COO格式 [源节点, 目标节点]
edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype=torch.long)
# 创建图数据对象
data = Data(x=x, edge_index=edge_index)

# 查看图的基本信息
print(f"节点数: {data.num_nodes}")
print(f"边数: {data.num_edges}")
print(f"节点特征维度: {data.num_node_features}")

GNN核心原理简析

图神经网络如何处理图数据？GNN通过消息传递机制实现节点特征的更新，核心思想是：每个节点的特征会受到其邻居节点特征的影响。具体来说，GNN层会聚合每个节点的邻居特征，并结合自身特征进行更新，从而实现信息在图中的传播。

场景应用

社交网络分析

社交网络中的用户关系如何用图神经网络建模？在社交网络中，用户可以表示为节点，用户之间的关注关系表示为边。GNN可以用于用户兴趣预测、社区发现等任务。例如，通过分析用户的社交关系和行为特征，预测用户可能感兴趣的内容。

分子结构预测

如何利用GNN预测分子属性？分子可以自然地表示为图结构，其中原子是节点，化学键是边。GNN能够学习分子的结构特征，从而预测分子的性质，如药物分子的活性、毒性等。这在药物发现和材料科学中具有重要应用。

实践操作

环境搭建与配置

⏱️ 预计耗时：5分钟

如何快速安装PyG并验证环境？推荐使用pip安装PyG：

# 基础安装
pip install torch_geometric

# 如需完整功能，可从源码安装
git clone https://gitcode.com/GitHub_Trending/py/pytorch_geometric
cd pytorch_geometric
pip install -e .[full]

安装完成后，可以运行示例文件验证环境是否配置成功：

python examples/cora.py

数据加载与预处理

⏱️ 预计耗时：10分钟

如何加载和处理图数据集？PyG内置了多种图数据集，以CiteSeer学术论文数据集为例：

from torch_geometric.datasets import Planetoid

# 加载CiteSeer数据集
dataset = Planetoid(root='data/CiteSeer', name='CiteSeer')
data = dataset[0]  # 获取图数据对象

# 查看数据集信息
print(f"数据集名称: {dataset.name}")
print(f"类别数: {dataset.num_classes}")
print(f"节点特征维度: {dataset.num_node_features}")
print(f"训练集大小: {data.train_mask.sum().item()}")
print(f"验证集大小: {data.val_mask.sum().item()}")
print(f"测试集大小: {data.test_mask.sum().item()}")

对于大规模图数据，可以使用NeighborLoader进行高效采样：

from torch_geometric.loader import NeighborLoader

# 创建邻居采样加载器
loader = NeighborLoader(
    data,
    num_neighbors=[10, 5],  # 每层采样的邻居数
    batch_size=32,          # 批次大小
    input_nodes=data.train_mask,  # 训练节点
)

# 迭代加载数据
for batch in loader:
    print(f"批次节点数: {batch.num_nodes}")
    print(f"批次边数: {batch.num_edges}")
    break

图神经网络模型构建与训练

⏱️ 预计耗时：15分钟

如何构建一个简单的GNN模型？下面以GCN（图卷积网络）为例，实现一个节点分类模型：

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class GCN(torch.nn.Module):
    def __init__(self, hidden_channels, num_node_features, num_classes):
        super().__init__()
        # 第一层GCN卷积
        self.conv1 = GCNConv(num_node_features, hidden_channels)
        # 第二层GCN卷积，输出类别数
        self.conv2 = GCNConv(hidden_channels, num_classes)

    def forward(self, x, edge_index):
        # 第一层卷积 + ReLU激活
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=0.5, training=self.training)  # Dropout防止过拟合
        # 第二层卷积
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)  # 对数softmax输出

# 初始化模型
model = GCN(hidden_channels=16, num_node_features=dataset.num_node_features, num_classes=dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
criterion = torch.nn.NLLLoss()  # 负对数似然损失

# 训练函数
def train():
    model.train()
    optimizer.zero_grad()  # 清零梯度
    out = model(data.x, data.edge_index)  # 前向传播
    loss = criterion(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)  # 预测类别
    # 计算各数据集准确率
    train_acc = int((pred[data.train_mask] == data.y[data.train_mask]).sum()) / int(data.train_mask.sum())
    val_acc = int((pred[data.val_mask] == data.y[data.val_mask]).sum()) / int(data.val_mask.sum())
    test_acc = int((pred[data.test_mask] == data.y[data.test_mask]).sum()) / int(data.test_mask.sum())
    return train_acc, val_acc, test_acc

# 训练模型
for epoch in range(1, 201):
    loss = train()
    train_acc, val_acc, test_acc = test()
    if epoch % 10 == 0:
        print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Train: {train_acc:.4f}, Val: {val_acc:.4f}, Test: {test_acc:.4f}')

模型评估与可视化

⏱️ 预计耗时：5分钟

如何评估模型性能并可视化结果？可以使用PyG提供的工具进行模型评估和结果可视化：

import matplotlib.pyplot as plt
from sklearn.manifold import TSNE

# 获取模型输出
model.eval()
out = model(data.x, data.edge_index)

# 使用TSNE降维可视化节点嵌入
tsne = TSNE(n_components=2)
z = tsne.fit_transform(out.detach().numpy())

# 绘制节点嵌入
plt.figure(figsize=(10, 10))
plt.scatter(z[:, 0], z[:, 1], c=data.y, cmap='tab10')
plt.colorbar()
plt.title('Node Embeddings Visualization')
plt.show()

上图展示了GraphGPS混合模型架构，它结合了MPNN（消息传递神经网络）和Transformer的优势，能够有效捕捉图的局部和全局特征。

进阶拓展

自定义图神经网络层

如何根据需求自定义GNN层？PyG提供了灵活的接口，允许用户自定义GNN层。下面是一个简单的自定义GNN层示例：

import torch
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree

class CustomGNNLayer(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)
        
        # 消息传递
        return self.propagate(edge_index, x=x)

    def message(self, x_j, edge_index):
        # 计算归一化系数
        row, col = edge_index
        deg = degree(col, x_j.size(0), dtype=x_j.dtype)
        deg_inv_sqrt = deg.pow(-0.5)
        norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]
        
        # 返回归一化后的邻居特征
        return norm.view(-1, 1) * x_j

处理大规模图数据

如何处理无法完全加载到内存的大规模图数据？PyG提供了多种方法处理大规模图数据，如使用NeighborSampler进行邻居采样，或使用分布式训练。下面是使用NeighborSampler的示例：

from torch_geometric.loader import NeighborSampler

# 创建邻居采样器
sampler = NeighborSampler(
    data.edge_index, 
    node_idx=None, 
    sizes=[10, 5],  # 每层采样的邻居数
    batch_size=32, 
    shuffle=True, 
    num_workers=4
)

# 迭代采样数据
for batch_size, n_id, adj in sampler:
    print(f'Batch size: {batch_size}')
    print(f'Nodes in batch: {n_id.shape[0]}')
    print(f'Adjacency matrix shape: {adj.shape}')
    break