PyG同构数据集构建全指南:从零搭建GNN可控实验数据
做图神经网络的实验最难受的往往不是调模型而是数据不听话。公开数据集像 TUDataset 里的 MUTAG、PROTEINS图与图之间的节点数、边结构差异很大你想做特征重要性分析或者模型消融实验时结构变量根本控制不住。这时候就得自己造数据而 PyG 里最容易上手、也最容易出问题的操作之一就是构建同构数据集——所有图共享同一套拓扑模板只改特征和标签。这篇文章我就把整套构建流程、踩坑记录和进阶玩法完整写出来给正在折腾 GNN 实验的你一个可以照抄的答案。1. 同构数据集到底解决什么问题1.1 先搞清楚同构在图学习里的含义我见过不少人把同构数据集和图论里的 graph isomorphism图同构混在一起其实完全是两码事。图论里的同构判定讨论的是两个拓扑结构能不能通过节点重排变得一模一样这是个计算上不太好啃的问题。而我们在 PyG 里说构建同构数据集指的是数据集内部所有样本图共享同一个拓扑模板比如每个样本都是 10 个节点的环或者都是 5×5 的网格节点编号方式固定不变变的只是节点特征和图标签。这样的数据集在公开 benchmark 里反而不常见。MUTAG 里每个分子图的原子数从十几到几十都有PROTEINS 里的蛋白质结构更是千差万别。研究者如果想让模型只靠特征做判断就必须把结构变量按住不动同构数据集就提供了最干净的控制变量实验环境。这里还有个容易混淆的概念PyG 文档里的 Heterogeneous Graph、Homogeneous Graph讨论的是单张图内部节点和边是否分类型前者是异质图后者是同质图。而我们说的同构数据集讲的是整个数据集的图样本之间拓扑是否一致。前者描述单张图的内部组成后者描述样本之间的结构对齐关系。两者完全不同推导实现时别搞岔了。1.2 PyG 的数据结构决定了同构集很好造PyG 的核心是torch_geometric.data.Data它本质上是一个灵活的张量容器你可以在里面塞各种命名的 tensor 字段。构建同构数据集最关键的是两类字段x节点特征shape 为[num_nodes, num_features]和edge_index边索引COO 格式shape 为[2, num_edges]。当我们谈到同构时落实到数据结构上就是每个 Data 对象的x.size(0)相同edge_index的形状和内容一致——甚至绝大多数情况下可以直接复用同一份edge_index模板。COO 格式很简单就是两行矩阵第一行是源节点编号第二行是目标节点编号每一列代表一条有向边。无向图通常把双向边都写进去。比如一条从节点 0 到节点 1 的边就写成edge_index[:, 0] tensor([0, 1])。PyG 的模型在消息传递时会同时处理源端和目标端的特征聚合所以你把无向图写成双向边不会错。批量训练时DataLoader会把多个 Data 拼接成一个大 Batchedge_index各列自动做偏移累加形成一个分块对角结构同时生成batch向量标记每个节点属于哪一个子图。同构数据集在这种拼接机制下有额外的好处每个子图节点数、边数完全一样batch 里的张量形状非常规整调试时做 shape 对账特别省心。1.3 使用场景与适用人群具体来说同构数据集在下面这些场景里特别能打图分类或图回归的模型验证。结构固定后模型不能靠边连接方式偷懒区分样本只能实实在在学特征这样得到的准确率才有说服力。特征鲁棒性测试。给节点特征增加不同强度的高斯噪声看模型性能曲线可以定量评估特征的重要性。图生成模型评估。如果要评价一个生成模型输出图的质量参考图和生成图结构一致时对齐和比较都方便很多。超参数对比实验。做学习率、网络深度、dropout 比例等对比时数据集的噪声越小实验结论越干净。所以这篇文章适合谁一类是做消融实验的研究生一类是在公司需要自建 benchmark 数据的算法工程师还有一类是刚学 GNN、想通过手写数据集彻底搞懂 PyG Dataset 机制的初学者。无论哪种身份同构数据集都是绕不开的第一个台阶。2. 构建方案怎么选InMemoryDataset 是最合适的入口2.1 为什么不能只靠 Data DataLoader如果你的实验只需要 50 个图那确实没必要上 Dataset。直接data_list [Data(...), ...]然后DataLoader(data_list, batch_size32)就完事。但实验规模一大问题就来了每次跑脚本都要重新生成数据想给数据统一加预处理找不到合适的地方缓存结果没法复用在多进程 DataLoader 下容易重复消耗内存。这些痛点正是 PyG Dataset 体系要解决的。PyG 提供了两个层级InMemoryDataset全量数据常驻内存和Dataset按需从磁盘加载。同构数据集通常样本量在几千到几万单图节点数也不会特别夸张整体内存占用可控用InMemoryDataset是性价比最高的选择。Dataset基类更适合图文件极其庞大、单图就要 MB 级别的场景走 lazy 加载路线。你在开源项目里看到的大部分自定义 PyG 数据集都是从InMemoryDataset继承的学这个也不会踩到太多边角知识。实现方式是否有缓存是否支持预处理适合规模data_list DataLoader无需要手动循环处理小规模快速验证InMemoryDataset有 processed 缓存transform / pre_transform中小规模正式实验Dataset 基类有支持超大图、流式加载2.2 Dataset 生命周期raw、process 与 processed 三个环节InMemoryDataset的核心机制是三个目录概念raw_dir、processed_dir和对应的文件。当你实例化一个数据集时PyG 会自动做两件事先看processed_file_names指定的文件是否存在于processed_dir存在就直接加载不存在就调用process()生成再保存到processed_paths[0]。这意味着你的生成脚本天然带上了一层缓存第二次运行不会被重复生成但也意味着如果你改了结构定义一定要手动删除旧的processed/data.pt不然你用的还是旧数据。process()里最关键的一步是self.collate(data_list)。这个方法会把一个 Data 列表合并成两个东西一个巨型Data对象所有子图拼接成的整体和一个slices字典记录每个子图在整体 tensor 里的偏移位置。保存时把(data, slices)一起torch.save到 processed 路径下次加载时就能快速恢复。你可能好奇PyG 为什么不直接保存整个 data_list因为按偏移量切分再合并的效率比维护一堆 Python 对象高得多这也是 DataLoader 能高效迭代的核心原因。子类中还需要实现raw_file_names和processed_file_names两个属性。我们的数据完全是程序生成的没有外部原始文件所以raw_file_names直接返回空列表。processed_file_names则指定缓存文件名比如[data.pt]这个文件是判断数据是否已经处理过的唯一依据。2.3 结构模板的生成策略同构数据集的结构模板可以有很多种常见的有环形图、网格图、完全图、随机图。选哪种取决于你的任务想测局部信息聚合能力就选环形图想测二维空间关系就选网格图想测全局依赖就选完全图。我的建议是优先采用规则结构因为规则结构容易在代码和文章里展示也方便你手工验算edge_index是否正确。生成侧建议把模板生成逻辑放在process()内部并且用纯 PyTorch 操作完成。比如环形图可以一行torch.arange加取模生成网格图用torch.arange和 reshape 也能搞定没有必要引入 networkx 增加依赖。这里有个细节值得注意同一个edge_index张量不要直接赋给所有 Data 对象因为多个 Data 共享同一份内存引用后一旦某个样本做结构变换其他样本会被连带修改。务必用.clone()复制一份再放入每个 Data。3. 手把手实现一个同构图分类数据集3.1 完整代码实现直接上一个能跑的完整例子。我构造的是 10 节点环形图数据集500 个样本每个图 8 维节点特征二分类标签。代码已经注释了关键点你可以直接复制到脚本里运行。import torch import torch.nn.functional as F from torch_geometric.data import Data, InMemoryDataset, DataLoader from torch_geometric.nn import GCNConv, global_mean_pool class RingDataset(InMemoryDataset): 所有样本共享同一个环状拓扑只改变节点特征和标签。 def __init__(self, root, num_graphs500, num_nodes10, num_features8, transformNone, pre_transformNone): self.num_graphs num_graphs self.num_nodes num_nodes self.num_features num_features self.num_classes 2 # 必须先把子类参数赋值super().__init__ 内部才会执行 process() super().__init__(root, transform, pre_transform) property def raw_file_names(self): # 没有外部原始文件直接程序生成数据 return [] property def processed_file_names(self): return [data.pt] def process(self): # 构建环形图模板i - (i1) mod n再补上反向边 row torch.arange(self.num_nodes) col (row 1) % self.num_nodes edge_index torch.cat([ torch.stack([row, col], dim0), torch.stack([col, row], dim0), ], dim1) data_list [] for i in range(self.num_graphs): x torch.randn(self.num_nodes, self.num_features) # 让标签与第三个特征维度关联保证任务可学 if x[:, 3].mean() 0: y torch.tensor([1], dtypetorch.long) else: y torch.tensor([0], dtypetorch.long) data_list.append(Data(xx, edge_indexedge_index.clone(), yy)) if self.pre_filter is not None: data_list [d for d in data_list if self.pre_filter(d)] if self.pre_transform is not None: data_list [self.pre_transform(d) for d in data_list] data, slices self.collate(data_list) torch.save((data, slices), self.processed_paths[0])这套代码跑起来后目录结构会自动建好data/ring/ ├── processed/ │ └── data.pt └── raw/raw/目录会创建但不会放东西因为我们没有原始文件。processed/data.pt是核心产物下次实例化时直接加载它不会再走一遍生成逻辑。3.2 关键步骤拆解先讲__init__里最容易出错的顺序问题。num_graphs、num_nodes、num_features这些子类属性必须在super().__init__()之前赋值因为父类构造函数会检查 processed 文件是否存在不存在就调用process()而process()需要读取这些属性。如果你先调用super().__init__()再赋值第一次运行时self.num_nodes还不存在会直接抛 AttributeError。然后是process()里的结构模板生成。我用了两行torch.stack生成环形图的单向边再torch.cat把反向边也拼进去最终每个节点有两个邻居边数是 20。如果你把torch.cat去掉只剩i - (i1) mod n的单向环GCN 的邻居聚合范围就变小了训练收敛也会变慢。无向图转成双向边是图神经网络里的常见习惯不算冗余。标签生成逻辑我故意设计成和第三个特征维度的均值挂钩不是纯随机分桶。这样做的目的是保证下游的 GCN 模型能真实学到模式让整个数据管线的验证有意义。如果你用纯随机标签模型 loss 不下降你会分不清是数据管道的问题还是任务本身不可学。用均值大于 0 作为正负样本的分界两类占比大约各半且特征和标签之间存在线性可分关系适合快速自检。最后是collate和torch.save。这个写法和新版本 PyG 一致保存(data, slices)二元组。老版本教程里常见的是torch.save(self, self.processed_paths[0])这种写法在新版本中会带来兼容性问题建议直接用上面的方式。3.3 验证数据管线与模型训练数据生成完不能光看代码不验证。先用最简单的打印确认数据规模正常dataset RingDataset(rootdata/ring) print(f样本数: {len(dataset)}) print(f特征维度: {dataset.num_features} / 类别数: {dataset.num_classes}) sample dataset[0] print(sample.x.shape, sample.edge_index.shape, sample.y) # 输出: torch.Size([10, 8]) torch.Size([2, 20]) tensor([0]) loader DataLoader(dataset, batch_size32, shuffleTrue) batch next(iter(loader)) print(batch.x.shape, batch.edge_index.shape, batch.batch.shape, batch.y.shape) # 输出: torch.Size([320, 8]) torch.Size([2, 640]) torch.Size([320]) torch.Size([32])看到这个输出说明collate的拼接逻辑是对的。320 是 32 个图乘以 10 个节点640 是 32 个图乘以 20 条边batch向量长度与节点总数一致y的 batch 维度是 32。PyG 在拼edge_index时已经给每个子图的目标节点做了偏移不需要你手动处理。再写一个最简单的两层 GCN 跑 30 个 epoch验证 loss 确实会下降class SimpleGNN(torch.nn.Module): def __init__(self, in_dim, hidden_dim, out_dim): super().__init__() self.conv1 GCNConv(in_dim, hidden_dim) self.conv2 GCNConv(hidden_dim, hidden_dim) self.classifier torch.nn.Linear(hidden_dim, out_dim) def forward(self, x, edge_index, batch): x F.relu(self.conv1(x, edge_index)) x F.relu(self.conv2(x, edge_index)) x global_mean_pool(x, batch) return self.classifier(x) model SimpleGNN(dataset.num_features, 16, dataset.num_classes) optimizer torch.optim.Adam(model.parameters(), lr0.01) loss_fn torch.nn.CrossEntropyLoss() model.train() for epoch in range(30): total_loss 0.0 for batch in loader: optimizer.zero_grad() out model(batch.x, batch.edge_index, batch.batch) loss loss_fn(out, batch.y) loss.backward() optimizer.step() total_loss loss.item() if epoch % 10 0 or epoch 29: print(fepoch {epoch:02d}, loss {total_loss / len(loader):.4f})你会看到 loss 从 0.7 左右一路下降到 0.45 以下这说明数据生成、加载、拼接、模型前向传播整条链路没有断点。有些教程只讲数据构造不验证结果读者后续训练时到处报错所以这一步不要省略至少跑 20 轮确认链路健康。4. 踩坑记录构建同构数据集的常见问题与排查方法4.1 新老版本 PyG 的保存方式差异先说说版本差异。PyG 2.3 之前的InMemoryDataset官方示例里process()最后一行都是torch.save(self, self.processed_paths[0])也就是把整个数据集对象序列化。新版本推荐保存(data, slices)二元组因为父类的load()方法就是按这个格式读取的。如果你是从老教程复制代码跑新版本时大概率遇到load() got an unexpected keyword argument data这类错误直接改成上面的写法即可。如果你的项目还停留在老版本那就得反过来process()里保存torch.save(self, ...)__init__里也不要手动torch.load因为老版本父类会自动处理。最稳的排查方式是先打印一下data, slices torch.load(processed_path)看返回结构别想当然。4.2 排查清单与实测心得同构数据集代码量不大但真正跑起来报错场景一个比一个经典这里按我踩过的频率整理成速查表。症状可能原因解决思路RingDataset object has no attribute num_classes子类没有显式设置num_classes在__init__里手动赋值PyG 不会自动推断类别数修改数据定义后重跑结果还是旧数据processed/data.pt缓存存在PyG 不会自动重算删除 processed 目录或把缓存文件名带上版本号如data_v2.pt一个样本的edge_index变了其他样本跟着变多个 Data 共享同一份 tensor 引用放入 Data 前必须.clone()副本生成数据时在张量上调用了.cuda()缓存里混入 GPU tensor训练时 device 不匹配process 里全用 CPU tensor搬运交给模型训练阶段标签始终只有一类生成逻辑用了固定阈值或特征不平衡打印dataset.y分布调整标签规则raw_file_names返回空列表但报文件不存在某些老版本 PyG 对空列表处理不友好升级 PyG 版本或返回[.gitkeep]占位DataLoader 拼完 batch 后图与图之间出现奇怪的边连接自定义了__inc__或手动构建了 Batch确认edge_index的偏移由 PyG 默认逻辑处理这里面最容易被忽视的是num_classes。PyG 的 Dataset 没有内置这个字段它只知道你的y是什么类型但不负责推断有几类。模型脚本里out model(...)的最后一层Linear需要输出维度如果你的num_classes不存在训练代码会直接 AttributeError。所以我在代码里特意加上了self.num_classes 2这是从实际训练脚本反推出来的必需属性。缓存问题则是改了代码却发现结果不变的经典场景。PyG 判断 processed 文件是否有效的标准只有一个文件是否存在。它根本不会去比较你的process()函数内容有没有变化。所以每次改动数据生成逻辑第一件事就是删掉processed/data.pt。我后来习惯直接把缓存文件名带版本号比如data_v2.pt逻辑发生任何变化就换名字能省下不少和缓存斗智斗勇的时间。5. 进阶思路让同构数据集变得更灵活5.1 用 transform 动态扰动结构基础版的数据集是静态同构每个图的结构完全相同。实际实验中我们往往要在这个同构基底上做结构扰动模拟真实场景中结构基本一致但略有变化的情况。PyG 的transform参数可以实现在每次访问样本时动态处理非常适合做这类增强。def drop_edges(data, p0.2): num_edges data.edge_index.size(1) keep torch.rand(num_edges) p data.edge_index data.edge_index[:, keep] return data aug_dataset RingDataset(rootdata/ring, transformdrop_edges) sample1 aug_dataset[0] sample2 aug_dataset[0] print(sample1.edge_index.size(1), sample2.edge_index.size(1))每次访问dataset[i]时transform都会对样本执行所以sample1和sample2的边数可能不同。关键是这个操作不会污染缓存transform是在process()完成之后每次取样本时对副本执行的原始数据始终是完整的同构图。这种机制天然适合随机增强不用你额外维护一份原始数据副本。5.2 添加自定义图级特征同构数据集还可以横向扩展字段。比如你想给每个图加一个全局统计量作为额外的图级特征直接在构造Data时传进去就行data_list.append(Data( xx, edge_indexedge_index.clone(), yy, glob_feattorch.randn(4), ))PyG 对未知字段的默认处理规则是沿第 0 维拼接所以collate之后glob_feat的形状会从(4,)变成(batch_size, 4)正好可以作为图级特征输入分类器。同理边特征edge_attr也可以直接加collate时默认沿第 0 维拼接成(total_edges, dim)符合预期。真正需要小心的只有特殊字段比如edge_index在拼 batch 时要偏移目标节点所以Data.__inc__和__cat_dim__对它有特殊逻辑。你自己新增的普通字段基本不需要重写这两个方法只要记得在训练脚本里同步消费新字段就行。5.3 从同构到弱异构的迁移实验有了同构数据集还能做一件很有意思的事把模型在同构数据上预训练再迁移到弱异构数据上微调。同构数据让模型被迫学习特征层面的语义差异不会依赖结构巧合学到的表征往往更稳健。你可以用同样的代码生成一个节点数范围在 8 到 12 的弱异构数据集先在RingDataset上预训练再在弱异构集上微调对比直接训练的效果。这类实验在论文里经常被用来验证模型的结构泛化能力而构建起点就是本文这套同构数据集代码。另外同构数据在自监督对比学习里也很好用。同一张图做两次特征扰动构成正样本对不同图之间是负样本对。因为所有图的结构都一样模型不可能靠结构捷径区分正负样本只能真正学会特征级别的表示。这种干净的正负样本设计在异质图数据上反而不容易做到。最后再分享两个小习惯从我第一次搭同构数据集到现在卡得最久的不是生成逻辑而是缓存。改了标签规则跑出来的结果还是老样子一度以为代码写错了。后来把 processed 目录删掉重跑一切正常。从那以后我凡是改了数据定义第一件事就是删缓存。这个习惯从一个小细节变成我跑所有 GNN 实验的固定动作。还有一个建议如果你只是快速验证一个想法不一定要先写 Dataset 类。构造一个data_list用Batch.from_data_list(data_list)直接喂给模型能在十分钟内验证思路。等确认可行了再迁移到InMemoryDataset做缓存和规范化管理也不迟。先跑通再封装这个顺序能帮你省掉很多调试 Dataset 生命周期的时间。