Planetoid 数据集 PyG 2.6.0 加载实战：Cora/CiteSeer/PubMed 3 种分割方式对比

Planetoid 数据集 PyG 2.6.0 加载实战：Cora/CiteSeer/PubMed 三种分割方式深度解析

引言：为什么需要关注数据分割策略？

在构建图神经网络模型时，许多开发者往往将注意力集中在模型架构设计和超参数调优上，却忽略了数据分割策略对模型性能的关键影响。实际上，不同的数据分割方式可能导致模型评估结果出现显著差异，这在学术研究和工业实践中都值得高度重视。

Planetoid数据集（包含Cora、CiteSeer和PubMed三个经典文献引用网络）作为图神经网络研究的基准测试集，其内置的多种分割方式各有特点：

• 'public'：固定分割方案，便于结果复现和横向比较
• 'full'：最大化利用标注数据，适合小规模数据集
• 'random'：随机分割方案，反映模型在数据分布变化时的鲁棒性

本文将深入解析PyG 2.6.0中Planetoid数据集的加载机制，通过对比实验揭示不同分割方式的特点，并给出针对不同研究目标的实用选择建议。

1. 环境准备与基础数据加载

1.1 安装与导入必要库

确保使用PyTorch Geometric 2.6.0及以上版本，以获得完整的分割策略支持：

pip install torch torch-geometric==2.6.0

基础加载代码示例：

import torch from torch_geometric.datasets import Planetoid from torch_geometric.transforms import NormalizeFeatures # 加载Cora数据集（默认使用'public'分割） dataset = Planetoid(root='./data', name='Cora', transform=NormalizeFeatures()) data = dataset[0] print(f"数据集: {dataset}") print(f"特征维度: {dataset.num_features}") print(f"类别数: {dataset.num_classes}") print(f"图结构: {data}")

1.2 数据集基础统计信息

三个数据集的关键指标对比：

注意：CiteSeer数据集中存在孤立节点，这在预处理时需要特殊处理

2. 三种分割方式的技术实现剖析

2.1 'public'固定分割方案

这是最常用的基准测试方案，其特点包括：

• 训练集：每类固定20个节点（Cora:140，CiteSeer:120，PubMed:60）
• 验证集：固定500个节点
• 测试集：固定1000个节点

加载代码示例：

def load_public_split(dataset_name): dataset = Planetoid(root='./data', name=dataset_name, split='public') data = dataset[0] print(f"\n{dataset_name} - public分割:") print(f"训练集节点数: {data.train_mask.sum().item()}") print(f"验证集节点数: {data.val_mask.sum().item()}") print(f"测试集节点数: {data.test_mask.sum().item()}") return data cora_data = load_public_split('Cora') citeseer_data = load_public_split('CiteSeer') pubmed_data = load_public_split('PubMed')

2.2 'full'最大化利用方案

当标注数据稀缺时，'full'方案可以最大化利用已有标注：

def load_full_split(dataset_name): dataset = Planetoid(root='./data', name=dataset_name, split='full') data = dataset[0] print(f"\n{dataset_name} - full分割:") print(f"训练集节点数: {data.train_mask.sum().item()}") print(f"验证集节点数: {data.val_mask.sum().item()}") print(f"测试集节点数: {data.test_mask.sum().item()}") # 验证集和测试集与public相同，训练集包含剩余所有节点 return data cora_full = load_full_split('Cora')

2.3 'random'随机分割方案

随机分割更适合评估模型鲁棒性：

def load_random_split(dataset_name, num_train=20, num_val=500, num_test=1000): dataset = Planetoid( root='./data', name=dataset_name, split='random', num_train_per_class=num_train, num_val=num_val, num_test=num_test ) data = dataset[0] print(f"\n{dataset_name} - random分割:") print(f"训练集节点数: {data.train_mask.sum().item()}") print(f"验证集节点数: {data.val_mask.sum().item()}") print(f"测试集节点数: {data.test_mask.sum().item()}") return data cora_random = load_random_split('Cora')

3. 分割方式对比与性能影响

3.1 节点分布对比分析

各数据集在不同分割方式下的节点分布：

Cora数据集分割对比

CiteSeer特殊注意事项

• 原始数据中存在孤立节点，随机分割时可能导致某些节点无法参与训练
• 建议预处理时检查data.contains_isolated_nodes()并做相应处理

3.2 模型性能对比实验

使用GCN模型在不同分割方案下的测试准确率对比（5次实验平均值）：

关键发现：full分割通常表现最佳（因训练数据更多），但public分割的结果更具可比性

4. 实战建议与高级技巧

4.1 如何选择合适的分割策略？

根据研究目标选择分割方式：

• 算法对比研究：使用'public'确保结果可复现
• 小样本学习：采用'full'最大化利用有限标注
• 鲁棒性测试：使用'random'进行多次交叉验证
• 迁移学习研究：组合多种分割方式评估泛化能力

4.2 自定义分割的高级实现

当内置分割不满足需求时，可自定义分割方案：

from torch_geometric.data import Data def custom_split(data, train_ratio=0.1, val_ratio=0.2): num_nodes = data.num_nodes perm = torch.randperm(num_nodes) train_size = int(num_nodes * train_ratio) val_size = int(num_nodes * val_ratio) data.train_mask = torch.zeros(num_nodes, dtype=torch.bool) data.val_mask = torch.zeros(num_nodes, dtype=torch.bool) data.test_mask = torch.zeros(num_nodes, dtype=torch.bool) data.train_mask[perm[:train_size]] = True data.val_mask[perm[train_size:train_size+val_size]] = True data.test_mask[perm[train_size+val_size:]] = True return data # 应用自定义分割 dataset = Planetoid(root='./data', name='Cora') data = dataset[0] data = custom_split(data)

4.3 处理数据不平衡问题

Planetoid数据集的类别分布通常不均衡，可采用以下策略：

from torch_geometric.utils import degree from torch_scatter import scatter_mean # 计算类别分布 class_counts = scatter_mean(torch.ones(data.num_nodes), data.y).cpu().numpy() # 采用类别平衡采样 def balanced_sample(data, samples_per_class=20): train_mask = torch.zeros(data.num_nodes, dtype=torch.bool) for c in range(dataset.num_classes): idx = (data.y == c).nonzero().view(-1) idx = idx[torch.randperm(idx.size(0))[:samples_per_class]] train_mask[idx] = True return train_mask

5. 跨框架数据加载对比

5.1 PyG与其他框架的加载方式差异

5.2 性能优化技巧

对于大规模图数据（如PubMed），可采用以下优化：

# 使用NeighborLoader进行采样 from torch_geometric.loader import NeighborLoader loader = NeighborLoader( data, num_neighbors=[25, 10], # 两阶采样 batch_size=32, input_nodes=data.train_mask ) # 启用PyG的稀疏矩阵优化 torch.sparse.check_sparse_tensor_invariants(False)

结语：从数据分割到研究可复现性

在实际项目中，我们发现许多论文难以复现的关键原因往往不是模型架构，而是数据分割策略的不透明。通过系统理解PyG中Planetoid数据集的分割机制，研究者可以：

1. 更准确地比较不同论文的结果
2. 根据实际需求选择最适合的分割策略
3. 设计更严谨的实验方案
4. 提高研究成果的可复现性

最后提醒：当使用Planetoid数据集发表研究成果时，务必明确说明采用的分割方式，这是保证研究可比较性的基本要求。