1. 从一次数据加载报错说起第一次在 PyTorch Geometric 里自定义Data类做小批量训练时我遇到过一个很典型的报错RuntimeError: The size of tensor a (5) must match the size of tensor b (4) at non-singleton dimension 0。当时排查了半天最后发现问题出在自己重写的__inc__和__cat_dim__上。这两个方法在 PyG 的官方文档里只有寥寥几行说明但只要你开始处理图结构之外的自定义属性比如节点级别的标签序列、边上的时间戳、全局的统计特征它们就一定会跳出来影响你。这篇内容就围绕这两个方法展开把它们的含义、作用、底层拼接逻辑、常见坑点以及实际项目里的写法讲透。适合已经用过 PyG 的Data和DataLoader、但对自己写__inc__/__cat_dim__还不太有把握的人。读完你应该能明白为什么 PyG 需要这两个方法、它们分别在什么时机被调用、返回值怎么影响最终 batch 的 tensor 形状以及怎么根据自己的数据结构写出正确的实现。先说结论性的定位__inc__决定的是“每个样本的某个属性在拼接后索引要偏移多少”__cat_dim__决定的是“这个属性沿哪个维度拼接”。前者服务于图结构里的索引类字段比如edge_index后者服务于所有需要被 batch 起来的字段。两者配合才能让一堆小图正确地合并成一张大图。2. PyG 的 batch 机制到底在做什么2.1 从单张图到一张大图PyG 的DataLoader和普通 PyTorch 的DataLoader有一个本质区别普通场景下一个 batch 是把 N 个样本在第 0 维堆叠成一个更大的 tensor而图数据不能这么简单堆叠因为每张图的节点数、边数都不一样。PyG 采用的是“对角拼接”思路——把 N 张小图拼成一张大图节点全部排在一起边也全部排在一起然后用batch向量记录每个节点属于哪张原图。举个例子一个 batch 里有 3 张图节点数分别是 4、3、5。拼接后的大图有 12 个节点batch向量就是[0,0,0,0,1,1,1,2,2,2,2,2]。原来第一张图的edge_index是[[0,1],[1,2]]第二张图是[[0,1],[1,0]]拼接后第二张图的节点索引必须整体加上第一张图的节点数 4变成[[4,5],[5,4]]。这个“加多少”就是__inc__要回答的问题。2.2 为什么不能自动推断你可能会想节点数不就是x.size(0)吗为什么还要我手动写__inc__原因在于 PyG 无法知道你自定义的某个字段到底是不是“索引类字段”。比如你有一个字段叫node_ref它存的是每个节点指向的另一组节点的编号那它在拼接时就需要偏移但如果它存的是每个节点的某种连续测量值那它就不需要偏移。这两种情况在 tensor 的形状上可能完全一样PyG 没法靠形状判断语义所以把决定权交给开发者。__cat_dim__也是同理。默认情况下 PyG 认为所有字段都沿第 0 维拼接但有些字段的语义维度不是第 0 维。最典型的就是edge_index它的形状是[2, num_edges]拼接时应该沿第 1 维边数维度拼而不是第 0 维。如果沿第 0 维拼两张图的edge_index会变成[4, num_edges]语义完全错乱。3.__inc__的含义与作用拆解3.1 方法签名与调用时机__inc__的定义形式是def __inc__(self, key, value, *args, **kwargs): ...它接收当前字段的键名key和值value返回一个整数或 tensor表示该字段在拼接时需要增加的偏移量。调用时机是在DataLoader把多个Data对象合并成Batch对象的过程中PyG 会遍历所有字段对每个字段调用一次__inc__。返回值有三种常见形态返回单个整数所有元素加同一个值、返回与value形状匹配的 tensor每个元素加不同的值、返回 0不偏移。理解这三种形态的适用场景是写好__inc__的关键。3.2 默认行为与源码逻辑PyG 的Data.__inc__默认实现大致是这样的逻辑如果key是edge_index返回当前图的节点数如果key是face返回节点数其他情况返回 0。这个默认实现覆盖了最基础的图结构字段但一旦你有自定义的索引字段就必须重写。源码里真正做拼接的地方在torch_geometric/data/collate.py的collate函数。它会先收集所有样本的某个字段值然后调用__inc__拿到偏移量对每个样本的值做累加偏移最后用torch.cat沿__cat_dim__指定的维度拼接。这个流程是理解两个方法协作关系的关键。3.3 索引偏移的实际计算假设一个 batch 有 3 张图节点数分别是[4, 3, 5]。对于edge_index字段PyG 会这样处理第 0 张图偏移 0edge_index不变第 1 张图偏移 4所有索引加 4第 2 张图偏移 743所有索引加 7这个累加过程是 PyG 自动完成的你只需要在__inc__里告诉它“当前这张图的节点数是多少”。所以对于edge_index__inc__返回self.num_nodes就够了。但如果你有多个不同类型的索引字段比如edge_index指向节点、hyperedge_index指向超边那就要分别返回对应的数量。3.4 什么时候需要返回 tensor 而非标量有一种情况标量不够用字段里存的是“每个元素指向不同集合的索引”。比如你有一个字段pair_index形状是[2, num_pairs]第一行指向节点集合第二行指向边集合。这时候两行需要加不同的偏移量就得返回一个形状为[2, 1]的 tensor让广播机制生效。def __inc__(self, key, value, *args, **kwargs): if key pair_index: return torch.tensor([[self.num_nodes], [self.num_edges]]) return super().__inc__(key, value, *args, **kwargs)这种写法在处理异构图或二分图时很常见因为不同行代表不同类型的实体偏移基准不同。4.__cat_dim__的含义与作用拆解4.1 方法签名与返回值__cat_dim__的定义形式是def __cat_dim__(self, key, value, *args, **kwargs): ...它返回一个整数表示该字段在拼接时沿哪个维度。默认实现里edge_index返回-1也就是最后一维其他字段返回 0。返回-1和返回1在二维 tensor 上等价但用-1更稳妥因为不依赖具体维度数。4.2 为什么edge_index要沿第 1 维拼edge_index的形状约定是[2, num_edges]第 0 维固定是 2源节点和目标节点第 1 维是边数。拼接多张图时我们希望边数累加而第 0 维保持 2 不变所以必须沿第 1 维拼。如果错误地沿第 0 维拼结果会变成[2*num_graphs, num_edges]不仅形状错了语义也完全崩了——原本表示“源/目标”的两行会被打乱成“图0源/图0目标/图1源/图1目标”。4.3 自定义字段的维度选择对于自定义字段选哪个维度取决于它的语义。判断标准很简单哪个维度是“样本间可累加”的就沿哪个维度拼。比如节点特征x形状[num_nodes, feat_dim]节点数可累加沿第 0 维拼边特征edge_attr形状[num_edges, feat_dim]边数可累加沿第 0 维拼全局特征u形状[1, feat_dim]每个图一个全局向量沿第 0 维拼后变成[num_graphs, feat_dim]序列特征seq形状[seq_len, num_nodes]如果每个图的序列长度相同、节点数不同那应该沿第 1 维拼最后这个例子比较绕但实际项目中会遇到。比如每个节点有一条时间序列序列长度固定为 10节点数不定。这时候seq的形状是[10, num_nodes]拼接时应该沿第 1 维让节点数累加序列长度保持 10。4.4 返回None的特殊含义__cat_dim__还可以返回None表示该字段不参与拼接而是被收集成一个 list。这在处理变长字段时很有用比如每个图的标签序列长度不同没法拼成规则 tensor就可以返回None让 PyG 把它们放进一个 list 里。def __cat_dim__(self, key, value, *args, **kwargs): if key variable_labels: return None return super().__cat_dim__(key, value, *args, **kwargs)拿到 batch 后batch.variable_labels就是一个 list每个元素对应一张图的原始值。这个技巧在序列标注、变长文本等场景下非常实用。5. 两个方法的协作流程与完整示例5.1 一次完整 batch 的字段处理流程把两个方法放在一起看PyG 处理一个字段的流程是收集 batch 内所有样本的该字段值得到一个 list对每个样本调用__inc__拿到偏移量对每个样本的值加上对应偏移量如果是索引字段调用__cat_dim__拿到拼接维度用torch.cat沿该维度拼接所有值这个流程里__inc__只对需要偏移的字段有意义__cat_dim__对所有字段都有意义。两者是正交的一个管“加多少”一个管“怎么拼”。5.2 一个带自定义索引字段的完整 Data 类下面这个例子模拟一个场景每个图有若干节点节点之间除了普通边还有“跨图引用边”指向另一个图的节点。这种字段需要自定义偏移逻辑。import torch from torch_geometric.data import Data class CrossGraphData(Data): def __init__(self, xNone, edge_indexNone, cross_indexNone, **kwargs): super().__init__(xx, edge_indexedge_index, **kwargs) self.cross_index cross_index def __inc__(self, key, value, *args, **kwargs): if key edge_index: return self.num_nodes if key cross_index: # cross_index 第一行指向本图节点第二行指向全局节点池 return torch.tensor([[self.num_nodes], [0]]) return super().__inc__(key, value, *args, **kwargs) def __cat_dim__(self, key, value, *args, **kwargs): if key in [edge_index, cross_index]: return -1 return super().__cat_dim__(key, value, *args, **kwargs)这里cross_index的第二行偏移量是 0因为它指向的是一个全局共享的节点池不随图拼接而偏移。第一行偏移num_nodes因为它指向本图节点拼接后需要加上前面所有图的节点数。5.3 验证拼接结果是否正确写完自定义方法后一定要手动验证一次拼接结果。最直接的方式是构造两个小样本手动算一遍期望的 batch 内容再和DataLoader的输出对比。from torch_geometric.loader import DataLoader d1 CrossGraphData( xtorch.randn(3, 4), edge_indextorch.tensor([[0, 1], [1, 2]]), cross_indextorch.tensor([[0, 2], [5, 7]]) ) d2 CrossGraphData( xtorch.randn(2, 4), edge_indextorch.tensor([[0], [1]]), cross_indextorch.tensor([[1], [3]]) ) loader DataLoader([d1, d2], batch_size2) batch next(iter(loader)) print(batch.edge_index) # 期望: [[0,1,3], [1,2,4]] 第二张图索引加3 print(batch.cross_index) # 期望: [[0,2,4], [5,7,3]] 第一行加3第二行不变 print(batch.batch) # 期望: [0,0,0,1,1]如果输出和期望一致说明两个方法写对了。如果不一致重点检查__inc__的返回值形状和__cat_dim__的维度选择。6. 常见问题与排查技巧实录6.1 报错速查表报错信息可能原因排查方向size mismatch at dim 0__cat_dim__返回了错误的维度检查该字段的语义维度是否与返回值一致index out of range__inc__返回 0 或偏移量不足确认索引字段是否被正确识别并偏移expected 2D tensor, got 1D__cat_dim__返回 None 导致 list 化检查是否误将需要拼接的字段返回 Nonecat(): invalid dimension返回的维度超出 tensor 维度范围用 -1 代替固定正数避免维度数变化时报错拼接后形状对但语义错__inc__返回了标量但需要 tensor检查多行索引字段是否需要逐行不同偏移6.2 三个容易踩的坑第一个坑忘记调用 super()。重写__inc__或__cat_dim__时如果只处理自定义字段而直接返回 0 或 0 维会破坏 PyG 对edge_index等内置字段的默认处理。正确做法是在分支末尾调用super().__inc__(...)或super().__cat_dim__(...)让内置逻辑继续生效。第二个坑__inc__返回了 tensor 但形状不对。当返回 tensor 时它的形状必须能和value广播。比如value形状是[2, num]返回[2, 1]可以广播返回[2]就不行。我见过有人返回[2, num]的完整偏移矩阵虽然能跑但浪费内存没必要。第三个坑把非索引字段也做了偏移。有些字段名字里带index但语义上不是索引比如index_feature存的是每个节点的某种特征值。这种字段如果被__inc__偏移了数值就全错了。判断标准是这个字段的值是否会被用来做 tensor 索引操作。会就需要偏移不会就不需要。6.3 调试技巧打印中间结果PyG 的 collate 过程不太容易直接打断点但可以在__inc__和__cat_dim__里加 print观察每个字段被调用时的 key 和 value 形状。def __inc__(self, key, value, *args, **kwargs): inc super().__inc__(key, value, *args, **kwargs) print(f[inc] key{key}, value_shape{value.shape if hasattr(value, shape) else type(value)}, inc{inc}) return inc跑一个小 batch看输出里哪些字段的 inc 是 0、哪些不是 0就能快速定位问题字段。这个技巧在排查“某个字段拼接后数值不对”时特别有效。7. 进阶场景与性能考量7.1 异构图的字段处理异构图里节点和边都有类型__inc__需要根据类型返回不同的偏移量。PyG 的HeteroData已经内置了这套逻辑但如果你在Data上模拟异构结构就需要自己维护每种类型的计数。常见做法是在__inc__里根据key的前缀判断类型返回对应的num_nodes或num_edges。7.2 大批量下的内存优化__inc__返回 tensor 时如果每个样本都新建一个 tensor在 batch 很大时会有额外的内存开销。优化方式是把偏移量缓存起来或者用标量加广播代替完整 tensor。比如[2, 1]的 tensor 比[2, num]省很多内存而广播机制会自动扩展效果一样。7.3 与__len__和num_nodes的关系__inc__里经常用到self.num_nodes这个属性是 PyG 根据x或edge_index自动推断的。如果你的图没有x也没有edge_indexnum_nodes可能返回 0 或报错。这时候需要在__init__里显式设置self.num_nodes或者在__inc__里用其他方式计算节点数。我一般建议在自定义Data里显式维护num_nodes避免推断失败。8. 我个人的几条实操建议第一先跑通默认行为再重写。PyG 的默认__inc__和__cat_dim__能覆盖大部分标准图结构场景。只有在确实遇到自定义字段拼接错误时才去重写。过早重写容易引入不必要的 bug。第二写完必做小样本验证。构造两个节点数不同的小图手动算一遍期望的 batch 内容和实际输出对比。这个习惯帮我省下了大量调试时间尤其是处理多行索引字段时。第三优先用-1而不是固定维度号。__cat_dim__返回-1表示最后一维不管 tensor 是二维还是三维都适用。返回1在二维时等价但换成三维就错了。用-1更稳。第四索引字段命名要清晰。把需要偏移的字段统一用_index后缀命名不需要偏移的用_feature或_attr。这样在__inc__里可以用key.endswith(_index)批量处理减少遗漏。这个命名约定在团队协作时尤其有用别人一看字段名就知道要不要偏移。第五变长字段果断返回 None。遇到长度不固定的字段不要硬凑成 tensor 拼接直接让__cat_dim__返回 NonePyG 会把它收集成 list。后续在模型里用torch.nn.utils.rnn.pad_sequence或手动处理比在 collate 阶段强行对齐要干净得多。这套机制刚接触时确实有点绕但理解之后会发现它的设计很合理把“索引偏移”和“拼接维度”这两个正交的关注点分开让开发者按字段粒度精确控制。实际项目里只要把这两个方法写对PyG 的 batch 机制就能处理绝大多数复杂的图数据结构。