PyTorch中高效处理一对多关系的不同尺寸张量运算方案

我有两个张量：第一个是1D张量（示例含3个值），第二个是2D张量，其第一维度为与第一个张量对应的ID，呈一对多关系（示例形状为(6,2)）。

示例代码如下：

# e.g. simple example of dot product
import torch

a = torch.tensor([2, 4, 3])
b = torch.tensor([[0, 2], [0, 3], [0, 1], [1, 4], [2, 3], [2, 1]]) # 1st column is the index to tensor a, 2nd column is the value

output = [(2*2)+(2*3)+(2*1),(4*4),(3*3)+(3*1)]
output = [12, 16, 12]

我当前的处理方式是统计b中每个ID的出现次数（如[3,1,2]），再用torch.split分组后循环处理。小张量下运行正常，但当张量规模达百万级、分组数过万时，运行速度极慢，求更优解决方案？

最优解决方案：向量化分组求和 #

循环处理在大规模张量场景下会产生极高的Python层开销，直接利用PyTorch的底层向量化操作可以彻底解决这个问题，核心是用scatter_add_实现高效的分组累加：

实现代码 #

import torch

a = torch.tensor([2, 4, 3])
b = torch.tensor([[0, 2], [0, 3], [0, 1], [1, 4], [2, 3], [2, 1]])

# 提取ID列和值列
ids = b[:, 0]
values = b[:, 1]

# 计算每个元素的加权贡献：a中对应ID的数值 * b中的值
contributions = a[ids] * values

# 初始化输出张量，按ID分组累加贡献值
output = torch.zeros_like(a)
output.scatter_add_(dim=0, index=ids, src=contributions)

print(output)  # 输出: tensor([12, 16, 12])

关键说明 #

• a[ids]是向量化索引操作，直接生成与values同形状的张量，无需循环遍历每个ID
• scatter_add_是PyTorch底层优化的算子，能高效地将同一ID的贡献值累加到输出张量的对应位置，时间复杂度为O(n)（n为b的行数）
• 整个过程完全避开Python循环，即使处理百万级数据、上万分组，性能也能保持线性提升

内容的提问来源于stack exchange，提问作者Chien Poon