PyTorch Geometric图神经网络开发指南：从基础到实践的45分钟入门教程

PyTorch Geometric（PyG）是基于PyTorch的图神经网络库，专为处理图结构数据设计。本指南将通过核心功能解析、场景化应用和进阶实践三个模块，帮助开发者快速掌握图神经网络的构建与应用，实现从数据建模到模型部署的全流程开发。无论是学术研究还是工业应用，PyG都能提供高效可靠的图深度学习解决方案。

一、核心功能解析

构建图数据结构：从张量到图对象

PyG采用Data对象统一表示图数据，包含节点特征、边关系等核心元素。这种结构化设计使图数据处理变得简单直观：

from torch_geometric.data import Data
import torch

# 节点特征矩阵 [num_nodes, num_features]
x = torch.tensor([[0.2, 0.5], [1.1, 0.3], [0.7, 0.9]], dtype=torch.float)
# 边索引 [2, num_edges]，COO格式存储
edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype=torch.long)
# 创建图数据对象
data = Data(x=x, edge_index=edge_index, y=torch.tensor([0, 1, 0]))

关键属性包括：x（节点特征）、edge_index（边连接关系）、y（标签）、edge_attr（边特征）等。通过data.num_nodes和data.num_edges可快速获取图的基本信息，这种设计极大简化了图数据的预处理流程。

实现图注意力机制：构建GAT模型

图注意力机制（GAT）——一种能让模型自动关注重要节点的神经网络结构，通过注意力权重计算实现节点间的信息传递。以下是使用PyG实现的两层GAT模型：

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

class GAT(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, heads=4):
        super().__init__()
        # 第一层GAT，多头注意力
        self.conv1 = GATConv(input_dim, hidden_dim, heads=heads)
        # 第二层GAT，单头输出
        self.conv2 = GATConv(hidden_dim * heads, output_dim, heads=1)
        
    def forward(self, x, edge_index):
        x = F.elu(self.conv1(x, edge_index))  # 激活函数
        x = F.dropout(x, p=0.5, training=self.training)  # 防止过拟合
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)  # 分类输出

GraphGPS混合模型架构展示了MPNN与Transformer的融合设计，体现了PyG模块化组件的灵活性

处理大规模图数据：分布式邻居采样

面对超大规模图数据，PyG提供NeighborLoader实现高效的邻居采样，通过局部邻居信息近似全局图计算：

from torch_geometric.loader import NeighborLoader

# 定义采样器，每层采样10和5个邻居
loader = NeighborLoader(
    data,
    num_neighbors=[10, 5],  # 两层采样策略
    batch_size=64,
    input_nodes=data.train_mask,  # 仅从训练集节点开始采样
)

# 迭代训练
for batch in loader:
    out = model(batch.x, batch.edge_index)
    loss = F.nll_loss(out[batch.train_mask], batch.y[batch.train_mask])

分布式图采样示意图展示了跨设备节点分配与局部计算的过程，实现大规模图的高效训练

二、场景化应用指南

分子性质预测：从SMILES到分子图

药物研发中，分子性质预测是关键任务。PyG可将SMILES分子表达式转换为图结构，实现端到端预测：

from torch_geometric.datasets import MoleculeNet
from torch_geometric.transforms import AddHydrogen, Compose

# 加载分子数据集，添加氢原子特征
dataset = MoleculeNet(root='data/qm9', name='QM9', 
                     transform=Compose([AddHydrogen()]))
data = dataset[0]  # 获取第一个分子图

# 构建分子图模型
from torch_geometric.nn import GINConv, global_add_pool

class MolGIN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = GINConv(torch.nn.Linear(11, 64))  # GIN卷积层
        self.fc = torch.nn.Linear(64, 1)  # 预测头
        
    def forward(self, x, edge_index, batch):
        x = self.conv(x, edge_index).relu()
        x = global_add_pool(x, batch)  # 图级池化
        return self.fc(x)

该模型可预测分子的能量、极性等物理化学性质，在药物发现和材料科学中具有重要应用价值。

3D点云分类：点云数据的图表示

将三维点云转换为图结构，通过图神经网络实现物体分类：

from torch_geometric.datasets import ModelNet
from torch_geometric.transforms import SamplePoints, KNNGraph

# 加载点云数据集，采样1024个点并构建KNN图
dataset = ModelNet(root='data/ModelNet10', name='10',
                  transform=Compose([SamplePoints(1024), KNNGraph(k=6)]))

# 构建点云分类模型
from torch_geometric.nn import PointConv

class PointGNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = PointConv(local_nn=torch.nn.Linear(3, 64))
        self.classifier = torch.nn.Linear(64, 10)
        
    def forward(self, x, edge_index, batch):
        x = self.conv(x, edge_index).relu()
        x = global_max_pool(x, batch)
        return self.classifier(x)

点云数据的采样、分组与特征提取流程展示了PyG在3D数据处理中的应用

三、进阶实践技巧

构建异构图模型：处理多类型节点关系

社交网络、知识图谱等场景常包含多种类型的节点和关系，PyG的HeteroData对象支持异构图建模：

from torch_geometric.data import HeteroData

# 创建异构图数据对象
hetero_data = HeteroData()
# 添加不同类型节点特征
hetero_data['user'].x = torch.randn(100, 16)  # 100个用户节点
hetero_data['item'].x = torch.randn(500, 8)   # 500个物品节点
# 添加用户-物品交互边
hetero_data['user', 'interacts', 'item'].edge_index = torch.tensor([
    [0, 0, 1, 1],  # 用户节点索引
    [0, 1, 1, 2]   # 物品节点索引
])

# 使用异构图卷积层
from torch_geometric.nn import HeteroConv, GCNConv

class HeteroGNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = HeteroConv({
            ('user', 'interacts', 'item'): GCNConv(16, 32),
            ('item', 'rev_interacts', 'user'): GCNConv(8, 32),
        })
        
    def forward(self, x_dict, edge_index_dict):
        return self.conv(x_dict, edge_index_dict)

模型解释与可视化：分析GNN决策过程

PyG提供模型解释工具，帮助理解GNN的决策依据：

from torch_geometric.explain import Explainer, GNNExplainer

# 初始化解释器
explainer = Explainer(
    model=model,
    algorithm=GNNExplainer(epochs=200),
    explanation_type='model',
    node_mask_type='attributes',
    edge_mask_type='object',
)

# 解释特定节点的预测
explanation = explainer(data.x, data.edge_index, index=10)
print(f"重要节点特征掩码: {explanation.node_mask}")
print(f"重要边掩码: {explanation.edge_mask}")

通过解释器可识别对预测结果贡献最大的节点特征和边连接，增强模型的可解释性和可信度。

四、资源导航

• 官方文档：docs/source/index.rst
• 示例代码库：examples/
• 单元测试集：test/

通过这些资源，开发者可以深入学习PyG的高级特性和最佳实践，加速图神经网络的开发与应用。无论是基础的节点分类任务，还是复杂的异构图学习，PyG都提供了简洁而强大的工具支持，助力开发者在图深度学习领域快速迭代创新。