如何快速在Numpy/PyTorch中找出值大于阈值的元素索引

你的双重循环实现逻辑没问题，但在处理稍大一点的矩阵时，Python层面的循环绝对会成为明显的性能瓶颈——毕竟你已经在用PyTorch了，完全可以利用它的张量并行计算能力来提速，根本没必要手动写循环。下面给你适配PyTorch和Numpy的两种高效实现方案：

PyTorch 原生实现（推荐，充分利用硬件加速） #

全程用张量操作完成，不管是CPU还是GPU运行，速度都会比你的原代码快几个数量级：

import torch

# 假设 abs_cosine 是 PyTorch 张量，threshold 是你设定的阈值
# 1. 生成掩码：排除对角线元素
diag_mask = ~torch.eye(abs_cosine.shape[0], dtype=torch.bool, device=abs_cosine.device)
# 2. 生成掩码：筛选大于等于阈值的元素
value_mask = abs_cosine >= threshold
# 3. 合并两个掩码：同时满足非对角线+阈值条件
combined_mask = diag_mask & value_mask
# 4. 获取所有符合条件的索引
rows, cols = torch.where(combined_mask)
# 5. 转换成你需要的元组列表格式（如果业务需要的话）
sim_vec = list(zip(rows.tolist(), cols.tolist()))

这个方案完全跳过了Python循环，所有计算都在底层的C/CUDA层面完成，能充分利用硬件的并行能力。而且如果你之前是把PyTorch张量转成Numpy再处理，直接用这个方案还能避免数据在CPU和GPU之间来回搬运的额外开销。

Numpy 适配方案（如果因迁移需求必须用Numpy） #

同样用向量化操作替代循环，比手动循环高效得多：

import numpy as np

# 假设 abs_cosine 是 Numpy 数组，threshold 是指定阈值
# 1. 生成非对角线掩码
diag_mask = ~np.eye(abs_cosine.shape[0], dtype=bool)
# 2. 生成阈值筛选掩码
value_mask = abs_cosine >= threshold
# 3. 合并掩码
combined_mask = diag_mask & value_mask
# 4. 获取符合条件的索引
rows, cols = np.where(combined_mask)
# 5. 转成元组列表
sim_vec = list(zip(rows.tolist(), cols.tolist()))

提速原理 #

• 向量化操作是由底层优化的C代码实现的，能利用CPU多线程或者GPU的大规模并行计算能力，而Python循环是单线程执行，每次循环都有额外的解释器开销。
• 避免了Python层和张量/数组底层之间的频繁数据交互，减少了不必要的性能损耗。

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