python - 在 Numpy/PyTorch 中快速查找值大于阈值的索引

Question

任务

给定一个 numpy 或 pytorch 矩阵，找到值大于给定阈值的单元格的索引。

我的实现

#abs_cosine is the matrix
#sim_vec is the wanted

sim_vec = []
for m in range(abs_cosine.shape[0]):
    for n in range(abs_cosine.shape[1]):
        # exclude diagonal cells
        if m != n and abs_cosine[m][n] >= threshold:
            sim_vec.append((m, n))

疑虑

速度。所有其他计算都建立在 Pytorch 上，使用 numpy 已经是一种妥协，因为它将计算从 GPU 转移到了 CPU。纯 python for 循环会使整个过程变得更糟(对于小数据集已经慢了 5 倍)。 我想知道我们是否可以在不调用任何 for 循环的情况下将整个计算转移到 Numpy(或 pytorch)？

我能想到的改进(但卡住了...)

bool_cosine = abs_cosine > threshold

返回 True 和 False 的 bool 矩阵。但是我找不到快速检索 True 单元格索引的方法。

Answer 1

以下是 PyTorch(完全在 GPU 上)

# abs_cosine should be a Tensor of shape (m, m)
mask = torch.ones(abs_cosine.size()[0])
mask = 1 - mask.diag()
sim_vec = torch.nonzero((abs_cosine >= threshold)*mask)

# sim_vec is a tensor of shape (?, 2) where the first column is the row index and second is the column index

以下在 numpy 中工作

mask = 1 - np.diag(np.ones(abs_cosine.shape[0]))
sim_vec = np.nonzero((abs_cosine >= 0.2)*mask)
# sim_vec is a 2-array tuple where the first array is the row index and the second array is column index