本文介绍如何在 PyTorch 中零循环、高效地将源张量按不规则一对多映射关系聚合到目标张量中，核心方法是利用 torch.Tensor.scatter_add_ 配合索引展开与值重复，适用于图神经网络、稀疏特征聚合等场景。
本文介绍如何在 pytorch 中零循环、高效地将源张量按不规则一对多映射关系聚合到目标张量中，核心方法是利用 `torch.tensor.scatter_add_` 配合索引展开与值重复，适用于图神经网络、稀疏特征聚合等场景。
在深度学习实践中，常需将一个较小的输入张量（如节点嵌入）依据不规则映射关系“分发”并累加到一个更大或结构不同的输出张量中（如超节点、聚合池化结果）。例如，每个 input[i] 可能贡献给多个 output[j]，而每个 output[j] 的值应为所有映射至它的 input[i] 之和。若用 Python 循环逐项更新，不仅性能低下，还无法被自动微分完整追踪——这正是 scatter_add 大显身手的典型场景。
核心思路：将“一对多”转化为“一维散列累加”
torch.Tensor.scatter_add_ 是 PyTorch 提供的原生、可导、GPU 加速的原子操作，其语义为：
对于每个 i，执行 out[index[i]] += src[i]（沿指定维度 dim）。
因此，我们只需将原始的一对多映射关系扁平化为两个对齐的一维张量：
- src：按映射频次重复 input 中的每个元素；
- index：将所有 mapping[i] 展开为连续索引序列；
- out：预分配的零初始化输出张量（长度 = max(index) + 1）。
完整实现示例
import torch # 输入定义 input = torch.tensor([0, 1, 2, 3], dtype=torch.float32) mapping = [[1], [0, 2, 4], [0, 3], [1, 2]] # 步骤 1：计算每个 input[i] 映射的目标数量（用于重复） reps = torch.tensor([len(x) for x in mapping]) # [1, 3, 2, 2] # 步骤 2：构建 src —— input[i] 重复 reps[i] 次 src = input.repeat_interleave(reps) # tensor([0, 1, 1, 1, 2, 2, 3, 3]) # 步骤 3：构建 index —— 所有目标索引按 mapping 顺序展开 index = torch.tensor([r for x in mapping for r in x]) # tensor([1, 0, 2, 4, 0, 3, 1, 2]) # 步骤 4：初始化输出张量（注意：size 必须覆盖最大索引） out_size = index.max().item() + 1 # → 5 out = torch.zeros(out_size, dtype=src.dtype) # 步骤 5：执行向量化累加（in-place） result = out.scatter_add(dim=0, index=index, src=src) print(result) # tensor([3., 3., 4., 2., 1.])
✅ 输出完全匹配预期：output[0] = input[1] + input[2] = 1 + 2 = 3，output[1] = input[0] + input[3] = 0 + 3 = 3，依此类推。
关键注意事项
- 索引安全性：index 中所有值必须满足 0 ≤ index[i] < out.size(dim)，否则触发 RuntimeError。建议显式校验：assert index.min() >= 0 and index.max() < out_size。
- 数据类型一致性：src 与 out 的 dtype 和 device 必须一致，否则报错；推荐统一使用 input.dtype 初始化 out。
- 内存效率：repeat_interleave 和列表推导式会创建中间张量，但全程无 Python 循环，全部在 CUDA 上并行执行，相比 for 循环提速 10–100×（取决于规模）。
- 梯度传播：scatter_add 是可导操作，反向传播会正确将梯度从 out 分配回 src（即原始 input），支持端到端训练。
- 扩展性提示：若 mapping 来自稀疏邻接表（如 COO 格式），可直接构造 index 和 src，无需 Python 层展开，进一步提升效率。
总结
scatter_add 是处理“一对多索引聚合”的黄金工具。它将原本需要嵌套循环或手动拼接的逻辑，压缩为三步张量变换 + 一次原语调用，兼具简洁性、高性能与可微性。掌握该模式，可显著优化 GNN 聚合、特征桶化、标签映射等任务的实现质量。记住口诀：展索引、复源值、零初化、scatter_add 一招制敌。
