在numpy/pytorch中，使用布尔数组或索引数组进行索引是否更快？

我可以用一个形状相同的布尔数组/张量或者一个包含我要查找的元素的整数索引的数组/张量来索引我的numpy数组/pytorch张量。哪个更快？

以下测试表明，使用numpy和pytorch中的索引数组，速度通常会快3到20倍：

In [1]: a = torch.arange(int(1e5))
idxs = torch.randint(len(a), (int(1e4),))
ind = torch.zeros_like(a, dtype=torch.uint8)
ind[idxs] = 1
ac, idxsc, indc = a.cuda(), idxs.cuda(), ind.cuda()

In [2]: %timeit a[idxs]
73.4 µs ± 1 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [3]: %timeit a[ind]
622 µs ± 8.99 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [4]: %timeit ac[idxsc]
9.51 µs ± 475 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [5]: %timeit ac[indc]
59.6 µs ± 313 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [6]: idxs = torch.arange(len(a)-1, dtype=torch.long)
ind = torch.zeros_like(a, dtype=torch.uint8)
ind[idxs] = 1
ac, idxsc, indc = a.cuda(), idxs.cuda(), ind.cuda()

In [7]: %timeit a[idxs]
146 µs ± 14.2 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [8]: %timeit a[ind]
4.59 ms ± 106 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)


In [9]: %timeit ac[idxsc]
33 µs ± 15.1 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [10]: %timeit ac[indc]
85.9 µs ± 56.9 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

---

正如先前的解所示，我期望基于整数的索引更快，因为输出张量的维数等于索引张量维数，这使得内存分配更容易。

在基于纯CPU的实现上考虑这种运行的值也很有趣。你期望得到类似的结果吗？

numpy

测试将使用dtype

bool

。我的经验是，布尔索引稍微慢一点，与第一次使用

np转换布尔值一致。非零

。更正，对于类似大小的问题，我得到10倍的速度差。这与对布尔值应用

非零

是一致的。@Denninger你是说pytorch只在CPU上运行？我希望它或多或少与numpy相同，因为pytorch以numpy格式存储数据，并可能使用numpy二进制文件进行计算（没有理由重新发明已经润滑良好的轮子：）。