如何在PyTorch中高效实现张量的一对多无循环映射

问题描述 #

假设我有一个张量input和一个不规则二维数组mapping，其中mapping[i]包含input[i]需要映射到output的索引列表。如何在不使用循环的情况下得到张量output，使得output[j]等于所有映射到它的input元素之和？input与output尺寸不同。

示例：

• input = [0,1,2,3]
• mapping = [[1], [0,2,4], [0,3], [1,2]]
• 对应output = [1+2, 0+3, 1+3, 2, 1] = [3,3,4,2,1]

解决方案（PyTorch版本） #

核心思路是将不规则映射展开为扁平的索引-元素对，再利用PyTorch的scatter_add_原生操作完成累加，全程无Python循环：

import torch

input = torch.tensor([0, 1, 2, 3])
mapping = [[1], [0, 2, 4], [0, 3], [1, 2]]

# 计算每个input元素需要重复的次数（对应mapping[i]的长度）
repeat_counts = torch.tensor([len(submap) for submap in mapping])
# 生成与映射索引一一对应的input元素序列
expanded_input = input.repeat_interleave(repeat_counts)
# 将二维mapping展平为一维目标索引数组
target_indices = torch.tensor([idx for sublist in mapping for idx in sublist])
# 确定output的长度（最大目标索引+1）
output_len = target_indices.max().item() + 1
# 初始化全0的output张量
output = torch.zeros(output_len, dtype=input.dtype)
# 按索引累加对应元素
output.scatter_add_(0, target_indices, expanded_input)

print(output)  # 输出: tensor([3, 3, 4, 2, 1])

解决方案（NumPy版本） #

如果使用NumPy，可通过np.add.at实现原地累加，同样无需循环：

import numpy as np

input = np.array([0, 1, 2, 3])
mapping = [[1], [0, 2, 4], [0, 3], [1, 2]]

# 展平映射索引和对应的input元素
target_indices = np.concatenate(mapping)
expanded_input = np.repeat(input, [len(submap) for submap in mapping])
# 确定output长度
output_len = target_indices.max() + 1
output = np.zeros(output_len, dtype=input.dtype)
# 按索引累加元素
np.add.at(output, target_indices, expanded_input)

print(output)  # 输出: [3 3 4 2 1]

原理说明 #

这种方法把一对多的映射关系转换成一对一的扁平序列，然后利用张量/数组库的底层优化操作完成累加，避免了Python循环带来的性能损耗，尤其适合处理大规模数据。

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