PyTorch几何深度学习扩展库实战指南
简介：PyTorch Geometric（PyG）是基于PyTorch的几何深度学习库，专为处理图结构数据设计，支持图神经网络（GNNs）的构建与训练。该库提供图数据结构封装、模块化图层、批处理支持、内置数据集和可视化工具等功能，广泛应用于社交网络、化学分子建模、推荐系统等领域。本资料介绍PyG的核心功能与使用流程，涵盖数据预处理、模型构建、训练评估等环节，适合Python与PyTorch开发者快速入门图神经网络开发。
1. 几何深度学习与图神经网络（GNN）概述
几何深度学习（Geometric Deep Learning）是深度学习领域的一个重要分支，旨在将传统深度学习方法从欧几里得空间（如图像、语音）拓展到非欧几里得结构数据（如图、流形）。图神经网络（Graph Neural Networks, GNN）作为其核心实现形式，能够有效处理节点与边构成的图结构数据，广泛应用于社交网络分析、推荐系统、分子结构预测等领域。
GNN的基本思想是通过消息传递机制（message passing），在图的节点之间传递和聚合信息，从而学习节点或图的嵌入表示。随着图结构的复杂性提升，GNN衍生出多种变体，如图卷积网络（GCN）、图注意力网络（GAT）、GraphSAGE等，以应对不同场景下的建模需求。
本章将为读者奠定GNN的理论基础，并引出后续章节中使用PyTorch Geometric（PyG）实现图神经网络的关键内容。
2. PyTorch Geometric库简介与安装
PyTorch Geometric（简称PyG）是一个基于PyTorch构建的图神经网络（GNN）开发框架，专为处理图结构数据而设计。它不仅提供了高效的图数据处理能力，还集成了大量经典的图神经网络模型模块与常用数据集，极大地简化了GNN模型的开发和实验流程。本章将从PyG库的核心功能、安装配置方式以及基本使用方法三个层面进行深入剖析，帮助读者快速上手并搭建起自己的图神经网络开发环境。
2.1 PyG库的核心功能与特点
作为专为图结构数据设计的深度学习库，PyTorch Geometric在功能设计和性能优化上具有鲜明的特点。其核心优势体现在对图结构数据的高效支持、与PyTorch生态系统的无缝集成，以及丰富的内置模块与数据集。
2.1.1 对图结构数据的高效支持
PyG针对图结构数据设计了专门的数据结构和处理机制，能够高效地进行图数据的存储、变换和批处理。图数据通常由节点（Node）和边（Edge）组成，PyG使用 Data 类来封装这些信息，具体包括：
-
x：节点特征矩阵，形状为[num_nodes, num_node_features] -
edge_index：边索引张量，形状为[2, num_edges]，采用COO格式存储 -
y：节点或图的标签 -
pos：节点的位置信息（用于三维图结构）
示例代码：创建一个简单的图结构
import torch
from torch_geometric.data import Data
# 节点特征矩阵：3个节点，每个节点有2个特征
x = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.float)
# 边索引：表示节点之间的连接关系
edge_index = torch.tensor([[0, 1, 1, 2],
[1, 0, 2, 1]], dtype=torch.long)
# 创建Data对象
data = Data(x=x, edge_index=edge_index)
print(data)
代码逻辑分析：
- 第1~2行导入PyG的核心模块；
-
x是一个包含3个节点的特征矩阵，每个节点有两个维度的特征； -
edge_index使用COO格式表示图中的边，第一行为源节点，第二行为目标节点； - 使用
Data()构造函数将图结构封装为一个数据对象； - 打印结果将显示图的结构信息，包括节点数、边数和特征维度。
表格：图数据结构字段说明
| 字段名 | 含义说明 | 数据类型 |
|---|---|---|
| x | 节点特征矩阵 | Tensor |
| edge_index | 边索引张量（COO格式） | LongTensor |
| y | 节点或图的标签 | Tensor |
| pos | 节点坐标（用于几何图） | Tensor |
| batch | 批处理索引（用于图批量处理） | LongTensor |
| edge_attr | 边的特征（可选） | Tensor |
通过这种结构化的封装方式，PyG能够实现对图数据的高效操作与批量处理，尤其适合在深度学习模型中进行大规模图结构训练。
2.1.2 与PyTorch生态系统的无缝集成
PyTorch Geometric构建在PyTorch之上，完全兼容PyTorch的张量操作、自动求导机制以及模型构建流程。这意味着开发者可以使用熟悉的PyTorch语法来定义图神经网络层、优化器、损失函数等组件，同时利用PyG提供的图处理模块来构建完整的GNN模型。
示例代码：在PyG中定义一个GCN模型
import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
class GCN(torch.nn.Module):
def __init__(self, num_features, hidden_dim, num_classes):
super(GCN, self).__init__()
self.conv1 = GCNConv(num_features, hidden_dim)
self.conv2 = GCNConv(hidden_dim, 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)
model = GCN(num_features=2, hidden_dim=16, num_classes=2)
print(model)
代码逻辑分析：
- 使用
GCNConv实现两个图卷积层； -
forward()方法中，图的特征x和边索引edge_index作为输入传递给卷积层； - 使用ReLU激活函数和Dropout层进行非线性变换；
- 最后通过
log_softmax输出分类结果； - 模型继承自
torch.nn.Module，与PyTorch模型构建方式一致。
Mermaid流程图：GCN模型的数据流动示意图
graph TD
A[输入图数据] --> B(GCNConv Layer 1)
B --> C[ReLU激活]
C --> D[Dropout]
D --> E(GCNConv Layer 2)
E --> F[LogSoftmax输出]
通过PyTorch与PyG的紧密结合，开发者可以在不改变原有PyTorch开发习惯的前提下，轻松实现图神经网络的构建与训练。
2.1.3 内置图神经网络模块与数据集
PyTorch Geometric内置了多种经典的图神经网络模块，如GCN、GAT、GraphSAGE、TopKPooling等，同时还提供了多个常用的图数据集（如Cora、Citeseer、PubMed、PPI等），开发者无需手动实现底层网络结构或数据预处理即可快速构建模型并进行实验。
表格：PyG内置图神经网络模块示例
| 模块名 | 功能说明 |
|---|---|
| GCNConv | 图卷积网络层 |
| GATConv | 图注意力网络层 |
| SAGEConv | GraphSAGE网络层 |
| EdgeConv | 边卷积网络层 |
| TopKPooling | 图粗化与图分类的池化操作 |
| global_mean_pool | 图级平均池化 |
示例代码：加载Cora数据集
from torch_geometric.datasets import Planetoid
dataset = Planetoid(root='/tmp/Cora', name='Cora')
print(dataset)
代码逻辑分析：
- 使用
Planetoid类加载Cora数据集； -
root参数指定数据集的存储路径； -
name参数指定数据集名称； - 打印结果将显示数据集的详细信息，包括节点数、类别数、划分方式等。
Mermaid流程图：数据集加载与模型训练流程
graph LR
A[加载数据集] --> B[预处理与划分]
B --> C[构建模型]
C --> D[训练模型]
D --> E[评估性能]
通过PyG内置的模块与数据集支持，开发者可以将更多精力集中在模型设计与优化上，而不是底层实现与数据处理。
2.2 PyG的安装与环境配置
为了顺利使用PyTorch Geometric，开发者需要完成其安装与环境配置。PyG的安装方式支持pip和conda，并且对PyTorch版本有一定的依赖要求。在安装过程中，可能会遇到一些常见问题，需要根据提示进行排查与解决。
2.2.1 基于pip和conda的安装方式
PyTorch Geometric可以通过pip或conda进行安装，推荐使用pip安装方式，因为conda支持的版本可能稍有滞后。
pip安装命令：
pip install torch_geometric
conda安装命令：
conda install -c pyg pytorch-geometric
安装依赖说明：
- PyTorch ≥ 1.8.0
- Python ≥ 3.7
- CUDA支持（可选，推荐）
2.2.2 安装常见问题与解决方法
问题1：安装失败，提示版本不兼容
原因 ：PyG依赖PyTorch的版本，如果本地PyTorch版本不匹配，可能导致安装失败。
解决方法 ：卸载旧版本PyTorch并安装兼容版本：
pip uninstall torch
pip install torch==1.13.1+cu117 -f https://download.pytorch.org/whl/torch_stable.html
问题2：无法导入PyG模块
原因 ：安装路径不正确，或者Python环境未切换到正确的虚拟环境。
解决方法 ：检查当前Python环境是否正确，使用以下命令查看已安装包：
pip list | grep torch-geometric
问题3：GPU加速未生效
原因 ：PyTorch未正确识别CUDA环境。
解决方法 ：确认是否安装了支持CUDA的PyTorch版本：
import torch
print(torch.cuda.is_available()) # 应返回True
2.2.3 验证安装与基础测试代码
安装完成后，可以通过以下代码验证PyG是否安装成功，并测试基本功能。
示例代码：验证PyG安装
import torch
from torch_geometric.data import Data
# 创建一个简单图结构
x = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.float)
edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype=torch.long)
data = Data(x=x, edge_index=edge_index)
print("节点特征：\n", data.x)
print("边索引：\n", data.edge_index)
输出结果：
节点特征：
tensor([[1., 2.],
[3., 4.],
[5., 6.]])
边索引：
tensor([[0, 1, 1, 2],
[1, 0, 2, 1]])
该测试代码验证了PyG的基本图结构构建功能，若输出正常，说明安装成功。
2.3 PyG的基本使用方式
掌握PyG的基本使用方式是构建图神经网络模型的前提。开发者需要熟悉常用模块的引入方式、图数据的构建与可视化方法，以及简单模型的搭建与运行流程。
2.3.1 引入常用模块与函数
PyG的模块组织结构清晰，开发者可以根据需要引入特定模块。以下是常用模块的引入方式：
import torch
import torch.nn.functional as F
from torch_geometric.data import Data
from torch_geometric.nn import GCNConv, global_mean_pool
from torch_geometric.datasets import Planetoid
from torch_geometric.utils import to_networkx
上述代码引入了数据结构、图卷积层、数据集、网络可视化等常用模块，开发者可根据具体任务灵活使用。
2.3.2 图数据的构建与可视化
PyG支持将图结构转换为NetworkX图对象，从而实现图的可视化。
示例代码：图结构可视化
import matplotlib.pyplot as plt
import networkx as nx
from torch_geometric.utils import to_networkx
# 转换为NetworkX图对象
G = to_networkx(data, to_undirected=True)
# 绘制图结构
plt.figure(figsize=(5, 5))
nx.draw(G, with_labels=True, node_color='lightblue')
plt.show()
代码逻辑分析：
- 使用
to_networkx()将PyG的Data对象转换为NetworkX图； - 使用
nx.draw()进行绘图； -
with_labels=True显示节点编号； - 可视化结果有助于开发者理解图的结构和连接关系。
2.3.3 简单的GNN模型搭建与运行
结合前面介绍的内容，我们可以搭建一个简单的GNN模型并运行训练流程。
示例代码：运行一个简单的GNN模型
from torch_geometric.nn import GCNConv
import torch.optim as optim
class SimpleGNN(torch.nn.Module):
def __init__(self, num_features, num_classes):
super(SimpleGNN, self).__init__()
self.conv = GCNConv(num_features, num_classes)
def forward(self, data):
x, edge_index = data.x, data.edge_index
x = self.conv(x, edge_index)
return F.log_softmax(x, dim=1)
# 初始化模型和优化器
model = SimpleGNN(num_features=2, num_classes=2)
optimizer = optim.Adam(model.parameters(), lr=0.01)
# 训练循环
model.train()
for epoch in range(100):
optimizer.zero_grad()
out = model(data)
loss = F.nll_loss(out, torch.tensor([0, 1, 0])) # 假设标签
loss.backward()
optimizer.step()
if epoch % 10 == 0:
print(f'Epoch {epoch}, Loss: {loss.item():.4f}')
代码逻辑分析：
- 定义一个单层GCN模型；
- 使用Adam优化器；
- 每个epoch进行一次前向传播、损失计算与反向传播；
- 输出每10个epoch的训练损失值；
- 模型最终会尝试对节点进行分类。
通过该示例，开发者可以初步掌握PyG中图神经网络的构建与训练流程，为进一步深入学习打下坚实基础。
以上内容完整展示了PyTorch Geometric库的核心功能、安装配置方式以及基本使用方法，涵盖了图结构数据的封装、模型构建、数据集加载、可视化与训练流程等关键环节，为后续章节的深入学习做好了充分铺垫。
3. 图数据结构封装（Data类）
图神经网络（GNN）处理的是非欧几里得结构的数据，这类数据以图的形式存在，由节点（顶点）和边组成。为了在PyTorch Geometric（PyG）中高效地处理这些图数据，PyG 提供了 Data 类，作为图数据的标准封装形式。 Data 类不仅能够统一表示图结构信息，还支持灵活的数据扩展与批处理机制，是构建 GNN 模型的核心基础。
本章将深入探讨 Data 类的设计原理、结构组成以及其在图数据处理中的核心作用，并通过自定义图数据集的构建方法和图数据增强策略，展示其在实际应用中的强大灵活性。
3.1 Data类的设计原理与结构组成
PyG 中的 Data 类是用于表示图数据的核心类，其设计目标是提供一个统一的接口，使得图结构可以被高效地处理和传递。 Data 对象本质上是一个字典，支持自动推导图的属性，如节点特征、边索引、标签等。
3.1.1 节点特征矩阵与边索引张量
Data 类中最关键的两个属性是：
-
x：节点特征矩阵，形状为[num_nodes, num_node_features] -
edge_index：边索引张量，形状为[2, num_edges]，表示图中每条边连接的两个节点索引
import torch
from torch_geometric.data import Data
# 示例：创建一个包含4个节点、3条边的图
x = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]], dtype=torch.float) # 节点特征
edge_index = torch.tensor([[0, 1, 2, 3], [1, 2, 3, 0]], dtype=torch.long) # 边索引
data = Data(x=x, edge_index=edge_index)
print(data)
代码逻辑分析：
-
x是一个二维张量，表示每个节点的特征向量。 -
edge_index是一个二维张量，第一行表示源节点，第二行表示目标节点，构成图的邻接关系。 -
Data类自动将这些属性封装为图对象，并提供访问方式如data.x、data.edge_index等。
参数说明：
| 参数名 | 类型 | 含义 |
|---|---|---|
x | Tensor | 节点特征矩阵 |
edge_index | LongTensor | 边连接关系，形状为 [2, num_edges] |
y | Tensor | 图或节点的标签 |
pos | Tensor | 节点的空间坐标（用于图可视化） |
3.1.2 辅助属性（标签、位置等）
除了基本的图结构信息外， Data 类还支持其他辅助属性，例如：
-
y：图或节点的标签，用于分类或回归任务 -
pos：节点的空间坐标，常用于图可视化 -
edge_attr：边特征，表示边的属性信息
# 添加标签和节点位置
data.y = torch.tensor([0], dtype=torch.long)
data.pos = torch.tensor([[0, 0], [1, 0], [1, 1], [0, 1]], dtype=torch.float)
print(data)
参数说明：
| 参数名 | 类型 | 含义 |
|---|---|---|
y | Tensor | 标签，可用于节点或图分类任务 |
pos | Tensor | 节点坐标，用于可视化 |
edge_attr | Tensor | 边特征，形状为 [num_edges, num_edge_features] |
3.1.3 图数据的批量处理机制
在训练 GNN 模型时，通常需要将多个图组成一个批量进行处理。PyG 提供了 DataLoader 类来实现图数据的批量加载，它会自动将多个 Data 对象合并为一个 Batch 对象，并保留原始的图结构信息。
from torch_geometric.data import DataLoader
# 创建两个图数据
data1 = Data(x=torch.randn(4, 2), edge_index=torch.tensor([[0,1,2,3],[1,2,3,0]]))
data2 = Data(x=torch.randn(3, 2), edge_index=torch.tensor([[0,1,2],[1,2,0]]))
loader = DataLoader([data1, data2], batch_size=2)
for batch in loader:
print(batch)
print("Number of graphs in batch:", batch.num_graphs)
代码逻辑分析：
-
DataLoader将多个Data实例打包成一个批次，每个图的属性仍然保留。 - 每个图的节点数和边数可以不同，
Batch会自动拼接x和edge_index并记录每个图的起始索引。
批处理流程图：
graph TD
A[原始图列表] --> B[DataLoader]
B --> C[Batch对象]
C --> D[统一节点特征矩阵]
C --> E[统一边索引张量]
C --> F[图索引映射]
3.2 自定义图数据集的构建
虽然 PyG 提供了大量内置图数据集（如 Cora、Citeseer），但在实际应用中，往往需要构建自定义图数据集以适应特定任务。PyG 提供了两种主要的数据集基类： InMemoryDataset 和 Dataset ，分别适用于内存型和磁盘型数据集。
3.2.1 继承 InMemoryDataset 和 Dataset 类
-
InMemoryDataset：适合小规模图数据，一次性加载到内存中。 -
Dataset：适合大规模图数据，按需从磁盘读取。
以下是一个继承 InMemoryDataset 构建自定义图数据集的示例：
from torch_geometric.data import InMemoryDataset
class MyDataset(InMemoryDataset):
def __init__(self, root, transform=None, pre_transform=None):
super(MyDataset, self).__init__(root, transform, pre_transform)
self.data, self.slices = torch.load(self.processed_paths[0])
@property
def raw_file_names(self):
return ['data.csv']
@property
def processed_file_names(self):
return ['processed_data.pt']
def download(self):
# 下载数据的逻辑
pass
def process(self):
# 处理数据并保存为 Data 对象
data_list = [create_data(i) for i in range(100)]
data, slices = self.collate(data_list)
torch.save((data, slices), self.processed_paths[0])
参数说明：
| 方法名 | 功能描述 |
|---|---|
__init__ | 初始化函数，加载处理后的数据 |
raw_file_names | 返回原始数据文件名列表 |
processed_file_names | 返回处理后的数据文件名列表 |
download | 下载原始数据 |
process | 处理原始数据并保存为 Data 对象 |
3.2.2 数据预处理与缓存机制
为了提升数据加载效率，PyG 支持将处理后的数据缓存到磁盘。当 process() 方法首次运行时，会生成 .pt 文件；后续加载时将直接读取缓存，避免重复处理。
数据处理流程图：
graph LR
A[原始数据] --> B{是否已处理?}
B -->|是| C[加载缓存数据]
B -->|否| D[执行process方法]
D --> E[保存为.pt文件]
E --> F[返回Data对象]
3.2.3 实现自定义图数据加载流程
构建自定义图数据集的关键在于实现 process() 方法，将原始数据转换为 Data 对象列表。以下是一个简单的实现示例：
def create_data(index):
x = torch.randn(5, 2) # 每个图5个节点，每个节点2个特征
edge_index = torch.randint(0, 5, (2, 8)) # 随机生成8条边
y = torch.tensor([index % 2], dtype=torch.long)
return Data(x=x, edge_index=edge_index, y=y)
代码逻辑分析：
-
create_data()函数为每个图生成随机节点特征和边连接关系。 -
Data对象包含节点特征、边索引和图标签。 - 在
process()方法中调用该函数生成多个图，并调用collate()合并为Batch对象。
3.3 图数据增强与变换操作
在深度学习中，数据增强是提升模型泛化能力的重要手段。PyG 提供了 Transform 类来对图数据进行各种变换操作，包括标准化、归一化、随机边删除等。
3.3.1 使用 Transform 类进行数据增强
PyG 的 transform 模块提供了一系列预定义的数据增强函数，如：
-
NormalizeFeatures()：标准化节点特征 -
RandomEdgeDropout()：随机删除边 -
ToSparseTensor()：转换为稀疏张量
from torch_geometric.transforms import NormalizeFeatures, RandomEdgeDropout
transform = NormalizeFeatures() # 标准化节点特征
data = transform(data)
print(data.x)
参数说明：
| 变换类名 | 功能描述 |
|---|---|
NormalizeFeatures() | 对节点特征进行标准化处理 |
RandomEdgeDropout(p=0.1) | 以概率 p 删除边 |
ToSparseTensor() | 将邻接矩阵转换为稀疏张量 |
3.3.2 标准化、归一化与随机边删除
标准化操作是图数据预处理中的常见步骤，通常用于消除特征量纲差异。PyG 提供了 NormalizeFeatures 类，其内部逻辑如下：
class NormalizeFeatures:
def __call__(self, data):
x = data.x
mean = x.mean(0, keepdim=True)
std = x.std(0, keepdim=True)
data.x = (x - mean) / (std + 1e-5)
return data
参数说明：
-
mean：节点特征的均值 -
std：节点特征的标准差 -
eps=1e-5：防止除零错误
3.3.3 动态图生成与采样策略
在训练 GNN 模型时，动态图生成与采样策略对于处理大规模图数据至关重要。PyG 提供了 NeighborSampler 和 RandomNodeSampler 等工具来实现图的动态采样。
from torch_geometric.data import NeighborSampler
sampler = NeighborSampler(data.edge_index, sizes=[10, 5], batch_size=32, shuffle=True)
参数说明：
| 参数名 | 类型 | 含义 |
|---|---|---|
sizes | List[int] | 每层邻居采样的节点数 |
batch_size | int | 每个训练批次的节点数 |
shuffle | bool | 是否在每个 epoch 打乱数据顺序 |
图采样流程图：
graph TD
A[原始图] --> B[NeighborSampler]
B --> C[按层采样邻居节点]
C --> D[生成子图]
D --> E[用于模型训练]
通过本章的介绍，我们深入了解了 PyG 中 Data 类的结构与功能，掌握了构建自定义图数据集的方法，并学习了图数据增强与动态采样的实现策略。这些知识为后续构建图神经网络模型打下了坚实基础。
4. 图卷积层实现（GCN、GAT、GraphSAGE等）
在图神经网络（GNN）中，图卷积层是实现图结构数据建模的核心组件。与传统的卷积神经网络（CNN）在网格数据上进行局部感知不同，图卷积层通过定义节点之间的信息传递机制，能够在非结构化的图数据上提取有效的特征表示。PyTorch Geometric（PyG）提供了多种图卷积层的实现，包括经典的图卷积网络（GCN）、图注意力网络（GAT）以及邻居采样策略（如 GraphSAGE），这些模块构成了构建 GNN 模型的基础。
本章将深入探讨几种主流图卷积层的实现原理、PyG 中的使用方法以及不同模型之间的性能差异，帮助开发者理解如何在实际项目中选择和使用这些图卷积层。
4.1 GCN图卷积层的原理与实现
4.1.1 消息传递机制与邻接矩阵计算
图卷积网络（Graph Convolutional Network, GCN）是一种经典的图神经网络模型，其核心思想是基于图的邻接矩阵和节点特征矩阵进行信息聚合。GCN 的消息传递机制可以表示为：
\mathbf{X}^{(l+1)} = \tilde{D}^{-\frac{1}{2}} \tilde{A} \tilde{D}^{-\frac{1}{2}} \mathbf{X}^{(l)} \mathbf{W}^{(l)}
其中：
- $\mathbf{X}^{(l)}$ 是第 $l$ 层的节点特征矩阵；
- $\tilde{A} = A + I$ 是加入自环的邻接矩阵；
- $\tilde{D}$ 是 $\tilde{A}$ 的度矩阵；
- $\mathbf{W}^{(l)}$ 是可学习的参数矩阵。
GCN 的核心操作是将每个节点的特征与其邻居节点的特征加权平均，并通过线性变换提取更高层次的特征表示。
图结构信息传递流程（mermaid）
graph TD
A[节点特征X] --> B[邻接矩阵A]
B --> C[构建归一化传播矩阵]
C --> D[与特征矩阵相乘]
D --> E[应用可学习参数W]
E --> F[输出下一层特征]
4.1.2 GCN层的参数设置与前向传播
在 PyG 中， GCNConv 类实现了 GCN 层。其构造函数如下：
import torch
from torch_geometric.nn import GCNConv
class GCN(torch.nn.Module):
def __init__(self, num_features, hidden_dim, num_classes):
super(GCN, self).__init__()
self.conv1 = GCNConv(num_features, hidden_dim)
self.conv2 = GCNConv(hidden_dim, num_classes)
def forward(self, data):
x, edge_index = data.x, data.edge_index
x = self.conv1(x, edge_index)
x = torch.relu(x)
x = self.conv2(x, edge_index)
return torch.log_softmax(x, dim=1)
代码解析：
-
GCNConv(num_features, hidden_dim)：定义输入维度和隐藏层维度； -
x = self.conv1(x, edge_index)：执行图卷积操作，输入包括节点特征x和边索引edge_index； -
torch.relu(x)：激活函数； -
log_softmax：用于分类任务的概率输出。
4.1.3 PyG中GCNConv的使用示例
from torch_geometric.datasets import Planetoid
import torch.optim as optim
# 加载Cora数据集
dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]
# 初始化模型
model = GCN(num_features=dataset.num_features, hidden_dim=16, num_classes=dataset.num_classes)
optimizer = optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
# 训练过程
model.train()
for epoch in range(200):
optimizer.zero_grad()
out = model(data)
loss = torch.nn.functional.nll_loss(out[data.train_mask], data.y[data.train_mask])
loss.backward()
optimizer.step()
参数说明：
-
data.train_mask：训练节点的掩码； -
data.y：节点的标签； -
nll_loss：负对数似然损失函数。
4.2 GAT图注意力网络实现
4.2.1 注意力机制在图结构中的应用
图注意力网络（Graph Attention Network, GAT）引入了注意力机制，使得模型在聚合邻居节点信息时能够动态分配权重，提升模型的表达能力和可解释性。
其核心公式为：
\alpha_{ij} = \text{softmax} j(e {ij}) \quad \text{where} \quad e_{ij} = a(\mathbf{W}\mathbf{h}_i, \mathbf{W}\mathbf{h}_j)
其中：
- $e_{ij}$ 是节点 $i$ 对节点 $j$ 的注意力系数；
- $a$ 是一个可学习的注意力函数；
- $\mathbf{W}$ 是线性变换矩阵。
4.2.2 GAT层的计算流程与参数配置
在 PyG 中， GATConv 实现了该机制。以下是一个基本的 GAT 模型定义：
from torch_geometric.nn import GATConv
class GAT(torch.nn.Module):
def __init__(self, num_features, hidden_dim, num_classes, heads=8):
super(GAT, self).__init__()
self.conv1 = GATConv(num_features, hidden_dim, heads=heads)
self.conv2 = GATConv(hidden_dim * heads, num_classes, heads=1)
def forward(self, data):
x, edge_index = data.x, data.edge_index
x = self.conv1(x, edge_index)
x = torch.relu(x)
x = self.conv2(x, edge_index)
return torch.log_softmax(x, dim=1)
参数说明：
-
heads=8：多头注意力的数量； -
hidden_dim * heads：由于多头拼接，第二层输入维度为hidden_dim * heads； -
heads=1：最后一层通常使用单头输出。
4.2.3 多头注意力机制的实现细节
多头注意力通过并行计算多个注意力头，提升模型的鲁棒性。在 PyG 中， GATConv 自动处理多头拼接和平均操作。
示例训练代码片段：
model = GAT(num_features=dataset.num_features, hidden_dim=8, num_classes=dataset.num_classes, heads=8)
optimizer = optim.Adam(model.parameters(), lr=0.005, weight_decay=1e-4)
# 同GCN的训练方式，仅替换模型部分
4.3 GraphSAGE与邻居采样策略
4.3.1 GraphSAGE算法的核心思想
GraphSAGE（Graph Sample and Aggregate）是一种支持大规模图训练的 GNN 模型，其核心思想是在训练过程中对邻居节点进行采样，从而避免计算所有邻居带来的高内存消耗。
GraphSAGE 提供了多种聚合函数（均值、LSTM、池化等）来聚合邻居信息。
4.3.2 NeighborSampler的使用方法
PyG 中的 NeighborSampler 支持邻居采样机制，适用于大规模图训练：
from torch_geometric.data import NeighborSampler
sampler = NeighborSampler(data.edge_index, sizes=[25, 10], node_idx=data.train_mask,
batch_size=1024, num_workers=12)
参数说明：
-
sizes=[25, 10]：表示每层采样的邻居数量； -
batch_size=1024：每批次处理的节点数； -
num_workers=12：并行加载数据的线程数。
4.3.3 大规模图数据的训练优化
GraphSAGE 的训练流程通常采用分层采样策略，以下是一个使用 DataLoader 的示例：
from torch_geometric.loader import DataLoader
train_loader = DataLoader(data, batch_size=32, shuffle=True)
表格：不同采样策略对比
| 方法 | 适用场景 | 优点 | 缺点 |
|---|---|---|---|
| 全图训练 | 小规模图 | 精度高，训练稳定 | 内存占用大 |
| NeighborSampler | 中大规模图 | 内存友好，支持批量训练 | 需要合理设置采样大小 |
| ClusterGCN | 超大规模图 | 支持图划分，分布式训练支持 | 实现复杂，依赖图划分 |
4.4 其他图卷积层的扩展与对比
4.4.1 TopKPooling、EdgeConv等模块
PyG 提供了多种图卷积层的扩展实现，如：
-
TopKPooling：图粗化操作，用于图分类任务； -
EdgeConv：基于边特征的图卷积，适用于点云等任务； -
GraphConv：基于图邻接的通用卷积层； -
SAGEConv：GraphSAGE 的图卷积实现。
示例：TopKPooling 使用
from torch_geometric.nn import TopKPooling, GCNConv
class TopKModel(torch.nn.Module):
def __init__(self, num_features):
super(TopKModel, self).__init__()
self.conv1 = GCNConv(num_features, 128)
self.pool1 = TopKPooling(128, ratio=0.8)
def forward(self, data):
x, edge_index, batch = data.x, data.edge_index, data.batch
x = self.conv1(x, edge_index)
x, edge_index, _, batch, _ = self.pool1(x, edge_index, batch=batch)
return x
4.4.2 不同卷积层的性能与适用场景
表格：图卷积层性能对比
| 层类型 | 是否支持邻居采样 | 是否支持注意力机制 | 是否支持图粗化 | 推荐应用场景 |
|---|---|---|---|---|
| GCNConv | 否 | 否 | 否 | 节点分类、小图训练 |
| GATConv | 否 | 是 | 否 | 可解释性要求高的任务 |
| SAGEConv | 是 | 否 | 否 | 大规模图训练 |
| EdgeConv | 否 | 否 | 否 | 点云、3D图数据 |
| TopKPooling | 否 | 否 | 是 | 图分类任务 |
总结建议：
- 小图 + 精度优先 ：选择
GCNConv或GATConv； - 大图 + 内存受限 ：使用
NeighborSampler+SAGEConv； - 图分类任务 ：使用
TopKPooling进行图粗化； - 点云/边特征任务 ：考虑
EdgeConv或GatedGraphConv。
后续章节将围绕图神经网络的完整训练流程展开，包括数据加载、模型构建、训练循环与评估方法等内容，帮助开发者构建完整的 GNN 工程实践能力。
5. 图神经网络模型构建流程
图神经网络（GNN）的模型构建流程与传统深度学习模型在结构上有相似之处，但也因图结构数据的特殊性而有所不同。本章将深入解析 GNN 模型从数据加载到模型训练、评估的完整流程，结合 PyTorch Geometric（PyG）框架，提供清晰、结构化的实现思路。
5.1 模型构建的基本流程
构建一个完整的图神经网络模型，通常包括三个关键阶段：数据加载与预处理、网络结构设计与层堆叠、模型编译与参数初始化。
5.1.1 数据加载与预处理
图数据的加载不同于传统图像或文本数据，通常需要加载图结构（邻接边列表）、节点特征矩阵、标签等信息。
在 PyG 中，数据通常以 Data 类对象进行封装。我们可以通过内置数据集（如 Cora、Citeseer）或自定义数据集实现加载。
from torch_geometric.datasets import Planetoid
# 加载Cora数据集
dataset = Planetoid(root='data/Cora', name='Cora')
data = dataset[0]
参数说明：
-
root：数据集存储路径。 -
name：指定数据集名称。 -
dataset[0]：返回第一个图数据对象，包含节点特征data.x、边索引data.edge_index、标签data.y等。
数据预处理操作：
- 标准化节点特征（如归一化处理）
- 添加自环边（如
AddSelfLoops） - 图数据增强（如随机删除边）
5.1.2 网络结构设计与层堆叠
GNN 模型通常由多个图卷积层堆叠而成，例如 GCN、GAT、GraphSAGE 等。以 GCN 为例：
import torch
from torch.nn import Linear
from torch_geometric.nn import GCNConv
class GCN(torch.nn.Module):
def __init__(self, num_features, hidden_dim, num_classes):
super(GCN, self).__init__()
self.conv1 = GCNConv(num_features, hidden_dim)
self.conv2 = GCNConv(hidden_dim, num_classes)
def forward(self, data):
x, edge_index = data.x, data.edge_index
x = self.conv1(x, edge_index)
x = torch.relu(x)
x = self.conv2(x, edge_index)
return torch.log_softmax(x, dim=1)
代码逻辑分析：
-
GCNConv：PyG 提供的图卷积层，支持邻接矩阵的高效传播。 -
forward：定义前向传播过程，图结构信息通过edge_index传递。 -
log_softmax：用于多分类任务的输出激活函数。
层堆叠策略：
- 浅层堆叠（2~3层）：适用于节点分类任务，避免过平滑问题。
- 深层堆叠（>5层）：可用于图级任务，结合残差连接、跳跃连接等技巧。
5.1.3 模型编译与参数初始化
在 PyG 中，模型的编译和训练流程与 PyTorch 一致，包括损失函数选择、优化器配置、学习率调度等。
model = GCN(num_features=dataset.num_features, hidden_dim=16, num_classes=dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
参数说明：
-
num_features：输入节点特征维度。 -
hidden_dim：隐藏层维度。 -
lr：学习率。 -
weight_decay：L2正则化系数。
初始化策略：
- He 初始化（ReLU 激活函数配合）
- Xavier 初始化（Sigmoid/Tanh 激活函数配合）
5.2 模型的训练流程设计
训练流程是模型构建的核心部分，包括损失函数选择、优化器设置、训练循环实现与日志记录。
5.2.1 损失函数的选择与配置
根据任务类型不同，损失函数也有所不同：
| 任务类型 | 损失函数 | PyTorch 实现 |
|---|---|---|
| 节点分类 | 交叉熵损失（CrossEntropyLoss） | torch.nn.CrossEntropyLoss() |
| 图分类 | 图级交叉熵或均方误差 | CrossEntropyLoss / MSELoss |
| 链接预测 | 二元交叉熵损失（BCELoss） | torch.nn.BCEWithLogitsLoss() |
示例（节点分类）：
criterion = torch.nn.CrossEntropyLoss()
5.2.2 优化器的设置与学习率调整
常用的优化器包括 Adam、SGD、RMSprop 等。PyG 中通常使用 Adam 作为默认优化器。
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.5)
参数说明：
-
step_size：每隔多少步调整学习率。 -
gamma：学习率衰减因子。
5.2.3 训练循环的实现与日志记录
一个完整的训练循环通常包括前向传播、损失计算、反向传播、参数更新、评估与日志记录。
for epoch in range(100):
model.train()
optimizer.zero_grad()
out = model(data)
loss = criterion(out[data.train_mask], data.y[data.train_mask])
loss.backward()
optimizer.step()
scheduler.step()
# 日志记录
if (epoch + 1) % 10 == 0:
model.eval()
pred = model(data).argmax(dim=1)
correct = (pred[data.val_mask] == data.y[data.val_mask]).sum()
acc = int(correct) / int(data.val_mask.sum())
print(f"Epoch {epoch+1:03d}, Loss: {loss.item():.4f}, Val Acc: {acc:.4f}")
关键步骤说明：
-
data.train_mask：仅训练节点参与损失计算。 -
loss.backward()：反向传播计算梯度。 -
optimizer.step()：更新模型参数。 -
model.eval()：切换评估模式，用于验证和测试。
可视化训练流程（mermaid流程图）：
graph TD
A[初始化模型] --> B[加载数据]
B --> C[定义损失函数和优化器]
C --> D[训练循环开始]
D --> E[前向传播]
E --> F[计算损失]
F --> G[反向传播]
G --> H[参数更新]
H --> I[学习率调整]
I --> J{是否达到最大训练轮次?}
J -- 否 --> D
J -- 是 --> K[结束训练]
5.3 模型评估与调试技巧
模型训练完成后，评估与调试是确保模型泛化能力的重要步骤。
5.3.1 准确率、F1分数等指标计算
对于分类任务，可以使用准确率、F1 分数、AUC 等指标进行评估。
from sklearn.metrics import classification_report
model.eval()
pred = model(data).argmax(dim=1)
print(classification_report(data.y[data.test_mask].cpu(), pred[data.test_mask].cpu()))
示例输出：
precision recall f1-score support
0 0.82 0.79 0.80 56
1 0.75 0.81 0.78 43
2 0.85 0.83 0.84 54
3 0.78 0.72 0.75 47
...
accuracy 0.79 271
macro avg 0.79 0.79 0.79 271
weighted avg 0.79 0.79 0.79 271
5.3.2 模型过拟合与欠拟合的诊断
诊断过拟合/欠拟合通常通过训练损失与验证损失的对比：
| 现象 | 表现 | 对策 |
|---|---|---|
| 过拟合 | 训练损失低，验证损失高 | 增加正则化、减少层数、早停法 |
| 欠拟合 | 训练和验证损失都高 | 增加模型复杂度、调整学习率 |
示例判断逻辑：
# 假设已有训练和验证损失历史
if train_loss < val_loss and val_loss > threshold:
print("模型过拟合")
elif train_loss > threshold and val_loss > threshold:
print("模型欠拟合")
else:
print("模型训练正常")
5.3.3 使用TensorBoard进行可视化监控
TensorBoard 可用于监控训练过程中的损失变化、学习率变化、准确率等。
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter()
for epoch in range(100):
...
writer.add_scalar('Loss/train', loss.item(), epoch)
writer.add_scalar('Accuracy/val', acc, epoch)
writer.close()
启动 TensorBoard：
tensorboard --logdir=runs
TensorBoard 可视化内容示例（mermaid图）：
graph LR
A[TensorBoard] --> B[Loss/train]
A --> C[Accuracy/val]
A --> D[Learning Rate]
B --> E[损失曲线]
C --> E
D --> E
通过上述流程，我们完整构建了一个图神经网络模型的训练与评估体系。在后续章节中，我们将深入探讨具体任务的实现方式，包括节点分类、图分类、链接预测等。
6. 内置数据集使用（Cora、Citeseer、PubMed等）
图神经网络的应用离不开高质量的图数据集。PyTorch Geometric（PyG）内置了多个经典的图数据集，如 Cora、Citeseer、PubMed 等，它们广泛用于节点分类、图分类和链接预测等任务。本章将详细介绍这些数据集的结构、加载方式、预处理方法以及在实际任务中的使用方式。
6.1 Cora、Citeseer与PubMed数据集介绍
6.1.1 数据集的图结构与类别分布
Cora、Citeseer 和 PubMed 是三个经典的引文网络数据集，常用于图神经网络中的节点分类任务。它们的共同特点是每个节点代表一篇论文，边代表论文之间的引用关系，节点特征是词袋（Bag-of-Words）向量，节点类别表示论文的主题类别。
| 数据集 | 节点数 | 边数 | 特征维度 | 类别数 |
|---|---|---|---|---|
| Cora | 2,708 | 5,429 | 1,433 | 7 |
| Citeseer | 3,327 | 4,732 | 3,703 | 6 |
| PubMed | 19,717 | 44,338 | 500 | 3 |
这些数据集在结构上具有稀疏性和小世界特性，适合用于图神经网络的研究与实验。
6.1.2 数据集的划分方式与用途
这些数据集通常采用固定的训练集、验证集和测试集划分方式，适用于半监督学习任务。例如，Cora 数据集通常划分为：
- 训练集：每个类别各取 20 个样本
- 验证集：500 个随机样本
- 测试集：1000 个随机样本
在 PyG 中，这些划分方式通过数据集的 data.train_mask 、 data.val_mask 和 data.test_mask 属性来表示。
6.1.3 PyG中数据集的自动下载与加载
PyG 提供了便捷的 API 来加载这些数据集。以下是一个加载 Cora 数据集的示例代码：
from torch_geometric.datasets import Planetoid
dataset = Planetoid(root='data/Cora', name='Cora')
print(dataset)
data = dataset[0]
print(data)
代码逻辑分析：
-
Planetoid是 PyG 提供的一个类，用于加载 Cora、Citeseer 和 PubMed 数据集。 -
root参数指定数据集的存储路径。 -
name参数用于指定具体的数据集名称。 -
dataset[0]返回一个Data对象，包含图的结构信息和节点特征、标签等。
参数说明：
-
root：本地存储路径，如果不存在则自动下载。 -
name：支持'Cora','Citeseer','Pubmed'。 - 返回的
Data对象包含以下属性： -
x：节点特征矩阵（Tensor） -
edge_index：边的索引列表（Tensor） -
y：节点标签（Tensor） -
train_mask,val_mask,test_mask：布尔掩码，用于划分数据集
6.2 数据集的加载与预处理
6.2.1 Dataset类与DataLoader的配合使用
在 PyG 中， Dataset 是一个抽象类，用于封装图数据集。加载内置数据集后，可以通过 DataLoader 进行批量加载。以下是一个使用 DataLoader 的示例：
from torch_geometric.loader import DataLoader
loader = DataLoader(dataset, batch_size=32, shuffle=True)
for batch in loader:
print(batch)
mermaid 流程图：
graph TD
A[Dataset对象] --> B[DataLoader初始化]
B --> C[迭代器遍历]
C --> D[返回批量数据]
代码逻辑分析：
-
DataLoader可以对图数据进行批量处理，特别适用于图分类任务。 -
batch_size控制每次返回的图数量。 -
shuffle=True表示在每个 epoch 开始时打乱数据顺序。
参数说明：
-
batch_size：每次返回的图数量。 -
shuffle：是否打乱数据顺序。 -
num_workers：用于并行加载数据的进程数。
6.2.2 数据划分与交叉验证策略
在实际训练中，固定划分的数据集可能无法充分评估模型的泛化能力。因此，可以采用 K 折交叉验证策略。以下是一个使用 RandomNodeSplit 实现数据划分的示例：
from torch_geometric.transforms import RandomNodeSplit
transform = RandomNodeSplit(num_val=0.1, num_test=0.2)
data = transform(dataset[0])
表格：RandomNodeSplit 参数说明
| 参数名 | 类型 | 描述 |
|---|---|---|
| num_val | float | 验证集所占比例 |
| num_test | float | 测试集所占比例 |
| split | str | 划分类型，如 'train_rest' 等 |
逻辑分析：
-
RandomNodeSplit是一个数据变换类，用于生成新的训练/验证/测试掩码。 - 通过设置比例参数，可以灵活控制划分比例。
6.2.3 图数据的标准化与特征工程
图数据的特征通常需要进行标准化处理，以提高模型的训练效果。以下是一个标准化节点特征的代码示例：
import torch
from torch_geometric.data import Data
# 假设 x 是原始特征矩阵
x = data.x
mean = x.mean(dim=0)
std = x.std(dim=0)
x = (x - mean) / (std + 1e-6)
new_data = Data(x=x, edge_index=data.edge_index, y=data.y,
train_mask=data.train_mask, val_mask=data.val_mask,
test_mask=data.test_mask)
逻辑分析：
- 对每个特征维度进行均值归一化（Z-score）。
-
mean和std分别计算每个特征的均值和标准差。 -
1e-6是防止除零的平滑项。
注意事项：
- 特征标准化应在训练集上计算，验证和测试集应使用相同的均值和标准差。
- 在 PyG 中可通过自定义
Transform实现更通用的特征工程。
6.3 数据集的实际应用案例
6.3.1 在节点分类任务中的使用
节点分类是图神经网络最常见的任务之一。以下是一个使用 GCN 模型在 Cora 数据集上进行节点分类的完整示例：
import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
class GCN(torch.nn.Module):
def __init__(self, num_features, hidden_dim, num_classes):
super(GCN, self).__init__()
self.conv1 = GCNConv(num_features, hidden_dim)
self.conv2 = GCNConv(hidden_dim, 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)
# 模型初始化与训练
model = GCN(dataset.num_features, 16, dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
def train():
model.train()
optimizer.zero_grad()
out = model(data)
loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
loss.backward()
optimizer.step()
for epoch in range(200):
train()
代码逻辑分析：
- 使用
GCNConv实现两层图卷积网络。 -
F.relu和F.dropout用于引入非线性与正则化。 -
F.log_softmax输出对数概率用于分类。
参数说明：
-
num_features：输入特征维度。 -
hidden_dim：隐藏层维度。 -
num_classes：类别数量。
6.3.2 图分类任务的数据集适配
虽然 Cora、Citeseer 和 PubMed 主要用于节点分类，但也可以通过图级聚合方法（如全局平均池化）将其适配为图分类任务。以下是一个使用 global_mean_pool 的示例：
from torch_geometric.nn import global_mean_pool
class GCNForGraphClassification(torch.nn.Module):
def __init__(self, num_features, hidden_dim, num_classes):
super().__init__()
self.conv1 = GCNConv(num_features, hidden_dim)
self.conv2 = GCNConv(hidden_dim, hidden_dim)
self.classifier = torch.nn.Linear(hidden_dim, num_classes)
def forward(self, data):
x, edge_index, batch = data.x, data.edge_index, data.batch
x = self.conv1(x, edge_index)
x = F.relu(x)
x = self.conv2(x, edge_index)
x = global_mean_pool(x, batch) # 图级池化
return self.classifier(x)
逻辑分析：
-
global_mean_pool将节点特征聚合为图级特征。 -
batch参数用于区分不同图的数据。
参数说明：
-
batch：一个一维张量，表示每个节点所属的图索引。
6.3.3 链接预测任务的构建方式
链接预测任务的目标是预测图中是否存在边。可以将 Cora 等数据集用于构建链接预测任务，方法如下：
- 构建正样本与负样本：
from torch_geometric.utils import negative_sampling
edge_index, _ = remove_self_loops(data.edge_index)
neg_edge_index = negative_sampling(edge_index=edge_index,
num_nodes=data.num_nodes,
num_neg_samples=edge_index.size(1))
- 定义模型与损失函数：
def link_prediction_loss(pos_out, neg_out):
return -torch.log(torch.sigmoid(pos_out) + 1e-15).mean() \
- torch.log(1 - torch.sigmoid(neg_out) + 1e-15).mean()
逻辑分析：
-
negative_sampling用于生成负样本。 - 使用交叉熵损失函数进行边预测。
参数说明：
-
num_neg_samples：生成的负样本数量。 -
sigmoid函数用于将输出转换为概率。
通过本章的学习，读者可以熟练掌握 PyG 中 Cora、Citeseer、PubMed 等经典数据集的使用方法，包括加载、预处理、标准化、以及在节点分类、图分类和链接预测任务中的实际应用。下一章将进一步深入讨论 PyG 在不同图任务中的具体实战案例。
7. PyG在节点分类、图分类、链接预测中的应用
7.1 节点分类任务实战
7.1.1 任务定义与数据准备
节点分类任务是指在图结构中，为每个节点预测其所属的类别标签。常见的应用场景包括社交网络中的用户兴趣分类、学术论文的学科分类等。
在PyTorch Geometric中，Cora、Citeseer等经典数据集已经被封装为 Planetoid 类，可以直接加载使用。以下是一个加载Cora数据集的示例代码：
from torch_geometric.datasets import Planetoid
dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]
print(data)
输出结果类似如下，展示了图的基本结构：
Data(x=[2708, 1433], edge_index=[2, 10556], y=[2708], train_mask=[2708], val_mask=[2708], test_mask=[2708])
其中：
- x ：节点特征矩阵（2708个节点，每个节点1433维特征）
- edge_index ：图的边索引，形状为 [2, E]
- y ：节点类别标签
- train_mask 、 val_mask 、 test_mask ：训练、验证、测试集的索引掩码
7.1.2 模型构建与训练流程
我们可以使用 GCNConv 构建一个简单的图卷积网络进行节点分类任务：
import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
class GCN(torch.nn.Module):
def __init__(self, num_features, hidden_dim, num_classes):
super(GCN, self).__init__()
self.conv1 = GCNConv(num_features, hidden_dim)
self.conv2 = GCNConv(hidden_dim, 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')
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)
model.train()
for epoch in range(200):
optimizer.zero_grad()
out = model(data)
loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
loss.backward()
optimizer.step()
7.1.3 准确率评估与结果分析
训练完成后，可以通过验证集和测试集评估模型性能：
model.eval()
_, pred = model(data).max(dim=1)
correct = float(pred[data.test_mask].eq(data.y[data.test_mask]).sum().item())
acc = correct / data.test_mask.sum().item()
print(f'Test Accuracy: {acc:.4f}')
输出示例：
Test Accuracy: 0.8150
7.2 图分类任务实践
7.2.1 图级特征聚合方法（Global Pooling）
图分类任务的目标是对整个图进行分类，例如分子属性预测、社交网络类型判断等。PyG提供了多种全局池化方法，如 global_mean_pool 、 global_max_pool 等。
以下是一个使用 global_mean_pool 的示例：
from torch_geometric.nn import global_mean_pool
class GraphClassifier(torch.nn.Module):
def __init__(self, num_features, hidden_dim, num_classes):
super(GraphClassifier, self).__init__()
self.conv1 = GCNConv(num_features, hidden_dim)
self.conv2 = GCNConv(hidden_dim, hidden_dim)
self.lin = torch.nn.Linear(hidden_dim, num_classes)
def forward(self, data):
x, edge_index, batch = data.x, data.edge_index, data.batch
x = F.relu(self.conv1(x, edge_index))
x = self.conv2(x, edge_index)
x = global_mean_pool(x, batch) # 聚合图级特征
x = self.lin(x)
return F.log_softmax(x, dim=1)
7.2.2 使用TopKPooling进行图粗化
对于需要更精细图结构处理的任务，可以使用 TopKPooling 模块实现图粗化，保留重要节点并压缩图结构。
from torch_geometric.nn import TopKPooling, GCNConv
import torch.nn as nn
class GNNWithPooling(torch.nn.Module):
def __init__(self, num_features, hidden_dim, num_classes):
super(GNNWithPooling, self).__init__()
self.conv1 = GCNConv(num_features, hidden_dim)
self.pool1 = TopKPooling(self.conv1, ratio=0.8)
self.conv2 = GCNConv(hidden_dim, hidden_dim)
self.pool2 = TopKPooling(self.conv2, ratio=0.8)
self.classifier = nn.Linear(hidden_dim, num_classes)
def forward(self, data):
x, edge_index, batch = data.x, data.edge_index, data.batch
x = F.relu(self.conv1(x, edge_index))
x, edge_index, _, batch, _, _ = self.pool1(x, edge_index, batch=batch)
x = F.relu(self.conv2(x, edge_index))
x = global_mean_pool(x, batch)
return F.log_softmax(self.classifier(x), dim=1)
7.2.3 分类模型训练与评估
训练流程与节点分类类似，主要区别在于数据集类型为图级分类数据，如 TUDataset ：
from torch_geometric.datasets import TUDataset
dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES')
loader = DataLoader(dataset, batch_size=32, shuffle=True)
model = GraphClassifier(dataset.num_features, 64, dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
for epoch in range(100):
for data in loader:
out = model(data)
loss = F.nll_loss(out, data.y)
loss.backward()
optimizer.step()
optimizer.zero_grad()
7.3 链接预测任务实现
7.3.1 任务定义与负样本生成
链接预测任务旨在预测图中未观察到的边是否存在。PyG提供了 negative_sampling 函数用于生成负样本。
from torch_geometric.utils import negative_sampling
edge_index = data.edge_index
neg_edge_index = negative_sampling(edge_index=edge_index, num_nodes=data.num_nodes)
7.3.2 模型设计与边预测实现
使用图嵌入的方式进行链接预测，例如使用 GCN 生成节点表示后，通过点积判断边的存在概率：
class LinkPredictor(torch.nn.Module):
def __init__(self, gnn_model):
super(LinkPredictor, self).__init__()
self.gnn = gnn_model
def forward(self, data, edge_index):
z = self.gnn(data)
return (z[edge_index[0]] * z[edge_index[1]]).sum(dim=1)
# 示例使用
model = GCN(dataset.num_features, 16, 16) # 输出节点嵌入
predictor = LinkPredictor(model)
# 正样本与负样本预测
pos_out = predictor(data, edge_index)
neg_out = predictor(data, neg_edge_index)
7.3.3 模型性能评估与可视化
可以使用AUC-ROC等指标评估链接预测性能：
from sklearn.metrics import roc_auc_score
def get_link_labels(pos_edge_index, neg_edge_index):
return torch.cat([
torch.ones(pos_edge_index.size(1)),
torch.zeros(neg_edge_index.size(1))
], dim=0)
labels = get_link_labels(edge_index, neg_edge_index)
preds = torch.cat([pos_out, neg_out], dim=0).sigmoid()
auc = roc_auc_score(labels.detach().numpy(), preds.detach().numpy())
print(f'AUC: {auc:.4f}')
输出示例：
AUC: 0.9235
7.4 综合案例：使用PyG构建端到端应用
7.4.1 构建完整的训练与推理流程
将模型训练与推理封装为模块化流程：
class GNNPipeline:
def __init__(self, model, dataset):
self.model = model
self.data = dataset[0]
self.optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
def train(self, epochs=200):
self.model.train()
for epoch in range(epochs):
self.optimizer.zero_grad()
out = self.model(self.data)
loss = F.nll_loss(out[self.data.train_mask], self.data.y[self.data.train_mask])
loss.backward()
self.optimizer.step()
def evaluate(self):
self.model.eval()
_, pred = self.model(self.data).max(dim=1)
correct = float(pred[self.data.test_mask].eq(self.data.y[self.data.test_mask]).sum().item())
acc = correct / self.data.test_mask.sum().item()
return acc
# 使用示例
pipeline = GNNPipeline(model, dataset)
pipeline.train()
acc = pipeline.evaluate()
print(f'Test Accuracy: {acc:.4f}')
7.4.2 模型部署与接口封装
可将训练好的模型保存为 .pt 文件，并提供API接口用于推理：
torch.save(model.state_dict(), 'gcn_model.pt')
# 加载模型
model.load_state_dict(torch.load('gcn_model.pt'))
model.eval()
结合Flask或FastAPI构建REST接口：
from flask import Flask, request, jsonify
app = Flask(__name__)
@app.route('/predict', methods=['POST'])
def predict():
data = request.json['data']
# 预处理图结构
# ...
with torch.no_grad():
out = model(data)
return jsonify({'result': out.tolist()})
7.4.3 应用场景分析与性能调优
针对实际应用场景（如社交网络推荐、化学分子分类），可通过以下方式进行调优：
- 使用更复杂的GNN层（如GAT、GraphSAGE）
- 引入图注意力机制提升模型表达能力
- 优化数据加载器与批处理策略
- 使用混合精度训练（AMP）加速训练过程
- 在GPU集群上进行分布式训练
通过以上方式，可以在实际项目中将PyG应用于多种图任务，并构建端到端的图神经网络应用系统。
简介：PyTorch Geometric（PyG）是基于PyTorch的几何深度学习库，专为处理图结构数据设计，支持图神经网络（GNNs）的构建与训练。该库提供图数据结构封装、模块化图层、批处理支持、内置数据集和可视化工具等功能，广泛应用于社交网络、化学分子建模、推荐系统等领域。本资料介绍PyG的核心功能与使用流程，涵盖数据预处理、模型构建、训练评估等环节，适合Python与PyTorch开发者快速入门图神经网络开发。
码道开发者社区，聚焦华为云码道 CodeArts 代码智能体，沉淀 Agent、Skill、鸿蒙开发实战内容，供开发者查阅资料、交流技术、分享工程实践
更多推荐
22
0- 0
已为社区贡献3条内容
所有评论(0)