PyTorch向量化查找二维张量匹配元素对应索引的优化方案

场景与需求 #

图匹配任务中需处理两类尺寸不同的(i,j)元组长整型张量，记新索引张量为a，旧索引张量为b，示例如下：

import torch
a = torch.LongTensor([[0,1], [0,2], [1,3], [2,4], [3,5]])
b = torch.LongTensor([[1,2], [1,3], [2,3], [3,5]])

需要实现两个目标：

• 生成等长于a的掩码，标记a中哪些元组已经存在于b中
• 获取所有匹配元组在b中对应的位置索引

原有逐行循环实现存在Python层开销大、大规模数据下运行极慢的问题，以下是两种无循环的向量化实现方案。

方案1：广播机制向量化实现（中小规模数据适用） #

利用PyTorch广播机制批量完成所有元组的相等判断，完全消除Python循环：

# 扩展维度后批量比对元组相等性，match_matrix形状为(len(a), len(b))
# 矩阵中(i,j)位置为True代表a[i]和b[j]完全相等
match_matrix = (a[:, None, :] == b[None, :, :]).all(dim=2)
# 按行判断是否存在匹配，直接得到存在性掩码
existing_edges_mask = match_matrix.any(dim=1).long()
# 提取匹配位置对应的b中索引
corresponding_old_idx = match_matrix[existing_edges_mask.bool()].nonzero(as_tuple=True)[1].tolist()

运行结果和原循环实现完全一致：

print(existing_edges_mask) # tensor([0, 0, 1, 0, 1])
print(corresponding_old_idx) # [1, 3]

注意：该方案会生成形状为(len(a), len(b))的中间矩阵，当a、b长度均超过10万时会占用过大显存，仅适合万级以内规模的数据

方案2：哈希编码线性复杂度实现（大规模数据适用） #

将二维元组编码为唯一的一维整数，把二维匹配转化为一维值匹配，时间和内存复杂度均为线性，可支持百万级边规模的快速计算：

# 计算编码基数，保证每个(i,j)元组映射为唯一整数，无哈希冲突
base = max(a.max(), b.max()) + 1
a_code = a[:, 0] * base + a[:, 1]
b_code = b[:, 0] * base + b[:, 1]

# 快速判断a中元素是否在b中存在
existing_edges_mask = torch.isin(a_code, b_code)
# 构建b中编码值到索引的映射，O(1)查询对应位置
b_code_to_idx = torch.zeros(b_code.max() + 1, dtype=torch.long)
b_code_to_idx[b_code] = torch.arange(b.shape[0])
corresponding_old_idx = b_code_to_idx[a_code[existing_edges_mask]].tolist()

existing_edges_mask = existing_edges_mask.long()

注意：如果b中存在重复元组，需要先对b做去重处理，否则索引映射会被最后一次出现的重复元组覆盖

性能对比 #

• 原循环实现：Python层逐行遍历，十万级边规模下运行时间可达数秒到数十秒
• 广播向量化实现：所有计算在Tensor算子层完成，无Python循环开销，比原循环快100~1000倍，适合中小规模数据
• 哈希编码实现：线性时间复杂度，内存开销极低，百万级边规模下可在毫秒级完成计算，是工业级大图场景的首选方案

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