Planetoid 数据集在 PyG 2.6.0 中的深度实践：数据分割策略对节点分类任务的影响

1. 引言：理解 Planetoid 数据集的核心价值

在当今图神经网络（GNN）研究领域，Planetoid 数据集（包含 Cora、CiteSeer 和 PubMed 三个子集）已成为评估模型性能的黄金标准。这些数据集不仅提供了真实的学术引用网络结构，还包含了丰富的文本特征和明确的分类标签，使其成为研究半监督学习的理想选择。

PyTorch Geometric（PyG）作为当前最流行的图深度学习框架之一，在其 2.6.0 版本中对 Planetoid 数据集的支持进行了多项优化。其中最关键的改进之一便是提供了多种数据分割模式（'public'、'full'、'geom-gcn' 和 'random'），这些模式直接影响模型的训练过程和最终性能表现。

为什么数据分割如此重要？

• 影响模型对图结构信息的利用程度
• 决定半监督学习中的标签传播效果
• 关系到模型评估的公平性和可重复性
• 不同分割策略会导致模型看到不同比例的标记/未标记节点

2. PyG 2.6.0 中 Planetoid 数据集的加载与初始化

2.1 基础加载方式

在 PyG 2.6.0 中加载 Planetoid 数据集非常简单，但其中的参数选择却大有讲究。以下是基础加载代码示例：

from torch_geometric.datasets import Planetoid # 加载Cora数据集，使用默认的public分割 dataset = Planetoid(root='/tmp/Cora', name='Cora', split='public') data = dataset[0] # 获取图数据对象 print(f'数据集: {dataset}') print(f'节点数量: {data.num_nodes}') print(f'边数量: {data.num_edges}') print(f'特征维度: {dataset.num_features}') print(f'类别数量: {dataset.num_classes}') print(f'训练节点数: {data.train_mask.sum().item()}') print(f'验证节点数: {data.val_mask.sum().item()}') print(f'测试节点数: {data.test_mask.sum().item()}')

2.2 数据分割参数详解

PyG 2.6.0 提供了四种分割策略，每种策略对应不同的应用场景：

关键参数说明：

• num_train_per_class: 仅在random模式下有效，控制每类的训练样本数
• num_val: 验证集大小
• num_test: 测试集大小

3. 四种分割模式的深度对比分析

3.1 数据统计对比

我们以 Cora 数据集为例，对比不同分割模式下的数据分布：

注意：geom-gcn 实际上提供了10种不同的分割方案，上表中显示的是每种方案的节点数

3.2 实现细节差异

不同分割模式在底层实现上有显著区别：

1. public模式

   • 使用原始 Planetoid 论文中的固定分割
   • 训练集非常小（Cora仅5%节点）
   • 验证集和测试集固定

2. full模式

   • 除验证和测试节点外，所有节点都用于训练
   • 充分利用了半监督学习中的未标记数据
   • 更适合现代GNN模型

3. geom-gcn模式

   • 提供10种不同的固定分割
   • 每次访问数据集时会轮换使用不同分割
   • 需要平均多次运行结果以获得稳定评估

4. random模式

   • 每次加载数据集时随机生成新分割
   • 可通过设置随机种子复现结果
   • 灵活性最高但需更多实验确保稳定性

3.3 对模型性能的影响

为了量化不同分割模式的影响，我们使用简单的GCN模型在Cora数据集上进行测试：

import torch import torch.nn.functional as F from torch_geometric.nn import GCNConv class GCN(torch.nn.Module): def __init__(self, hidden_channels): super().__init__() self.conv1 = GCNConv(dataset.num_features, hidden_channels) self.conv2 = GCNConv(hidden_channels, dataset.num_classes) def forward(self, x, edge_index): x = self.conv1(x, edge_index) x = F.relu(x) x = F.dropout(x, p=0.5, training=self.training) x = self.conv2(x, edge_index) return x model = GCN(hidden_channels=16) optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4) criterion = torch.nn.CrossEntropyLoss() 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 def test(): model.eval() out = model(data.x, data.edge_index) pred = out.argmax(dim=1) acc = (pred[data.test_mask] == data.y[data.test_mask]).sum() / data.test_mask.sum() return acc.item() for epoch in range(1, 201): loss = train() if epoch % 50 == 0: test_acc = test() print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Test Acc: {test_acc:.4f}')

在不同分割模式下得到的测试准确率对比：

4. 实战：构建完整的节点分类流程

4.1 数据准备与模型定义

首先，我们需要根据任务需求选择合适的分割模式。对于需要最大化利用数据的场景，推荐使用full模式：

# 加载数据 dataset = Planetoid(root='/tmp/Cora', name='Cora', split='full') data = dataset[0] # 定义GAT模型 from torch_geometric.nn import GATConv class GAT(torch.nn.Module): def __init__(self, hidden_channels, heads): super().__init__() self.conv1 = GATConv(dataset.num_features, hidden_channels, heads=heads) self.conv2 = GATConv(hidden_channels*heads, dataset.num_classes, heads=1) def forward(self, x, edge_index): x = F.dropout(x, p=0.6, training=self.training) x = self.conv1(x, edge_index) x = F.elu(x) x = F.dropout(x, p=0.6, training=self.training) x = self.conv2(x, edge_index) return x model = GAT(hidden_channels=8, heads=8)

4.2 训练与验证策略

针对不同分割模式，需要采用不同的训练策略：

对于public/full模式：

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 def validate(): model.eval() out = model(data.x, data.edge_index) pred = out.argmax(dim=1) val_acc = (pred[data.val_mask] == data.y[data.val_mask]).sum() / data.val_mask.sum() return val_acc.item()

对于geom-gcn模式：

# 需要运行多次取平均 test_accs = [] for run in range(10): dataset = Planetoid(root='/tmp/Cora', name='Cora', split='geom-gcn') data = dataset[0] # ...训练代码... test_accs.append(test()) print(f'平均测试准确率: {np.mean(test_accs):.4f} ± {np.std(test_accs):.4f}')

4.3 结果分析与可视化

训练完成后，我们可以使用TSNE对节点嵌入进行可视化：

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

5. 高级技巧与最佳实践

5.1 处理类别不平衡

Planetoid 数据集中的类别分布并不完全均衡。以 CiteSeer 为例：

解决方法：

1. 在损失函数中使用类别权重

class_counts = torch.bincount(data.y[data.train_mask]) class_weights = 1. / class_counts criterion = torch.nn.CrossEntropyLoss(weight=class_weights)

2. 在 random 分割时使用分层抽样

dataset = Planetoid(root='/tmp/CiteSeer', name='CiteSeer', split='random', num_train_per_class=20, num_val=500, num_test=1000)

5.2 超参数调优建议

基于不同分割模式的调优策略：

5.3 跨数据集性能对比

为了全面评估模型，应该在三个数据集上测试：

关键发现：

• PubMed 虽然节点最多，但因特征维度低且类别少，反而更容易获得高准确率
• CiteSeer 由于特征维度高且存在孤立节点，通常表现最差
• Cora 作为中等规模数据集，最适合初步实验和算法开发

6. 扩展应用与进阶思考

6.1 结合其他PyG功能

PyG 2.6.0 提供了许多可与 Planetoid 数据集结合使用的强大功能：

1. 数据增强

from torch_geometric.transforms import NormalizeFeatures, AddRandomWalkPE dataset = Planetoid(root='/tmp/Cora', name='Cora', transform=NormalizeFeatures()) dataset = Planetoid(root='/tmp/Cora', name='Cora', pre_transform=AddRandomWalkPE(walk_length=10))

2. 使用异构图神经网络

from torch_geometric.nn import HeteroConv, GATConv, GCNConv class HeteroGNN(torch.nn.Module): def __init__(self, hidden_channels): super().__init__() self.conv1 = HeteroConv({ ('paper', 'cites', 'paper'): GATConv(-1, hidden_channels), ('paper', 'rev_cites', 'paper'): GCNConv(-1, hidden_channels) }) # ...

6.2 迁移学习策略

Planetoid 数据集上学到的知识可以迁移到其他图数据集：

1. 特征提取器迁移

# 先在Cora上训练 cora_dataset = Planetoid(root='/tmp/Cora', name='Cora') # ...训练模型... # 冻结特征提取层 for param in model.conv1.parameters(): param.requires_grad = False # 微调最后层 pubmed_dataset = Planetoid(root='/tmp/PubMed', name='PubMed') # ...继续训练...

2. 图结构迁移学习

• 将在小图上学习的图结构模式应用于大图
• 特别适用于PubMed到更大医学文献网络的迁移

6.3 处理大规模图的技巧

当处理PubMed等较大规模数据集时：

1. 使用邻居采样

from torch_geometric.loader import NeighborLoader loader = NeighborLoader( data, num_neighbors=[30, 10], batch_size=256, input_nodes=data.train_mask )

2. 分布式训练

import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP dist.init_process_group('nccl') model = DDP(model.to(device), device_ids=[rank])

在实际项目中，选择合适的数据分割策略需要综合考虑实验目的、计算资源和评估需求。对于大多数研究场景，建议同时尝试public和full模式以获得全面认识；对于生产环境，random模式配合交叉验证可能更为可靠。