在 numpy/pytorch 中使用 bool 数组或索引数组进行索引是否更快？答案

【问题标题】：Is indexing with a bool array or index array faster in numpy/pytorch?在 numpy/pytorch 中使用 bool 数组或索引数组进行索引是否更快？
【发布时间】：2019-09-04 06:57:24
【问题描述】：

我可以使用相同形状的布尔数组/张量或包含我所追求的元素的整数索引的数组/张量来索引我的 numpy 数组/pytorch 张量。哪个更快？

【问题讨论】：

标签： numpy pytorch numpy-slicing

【解决方案1】：

以下测试表明，在 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 的实现上考虑此类运行的值也会很有趣。你期望那里有类似的结果吗？
A numpy 测试将使用 dtype bool。我的经验是布尔索引稍慢，与首先使用 np.nonzero 转换布尔值一致。
更正，对于类似大小的问题，我得到 10 倍的速度差异。这与将nonzero 应用于布尔值是一致的。
@dennlinger 你的意思是 pytorch 只在 CPU 上运行？我希望它或多或少与 numpy 相同，因为 pytorch 以 numpy 格式存储数据，并且可能也使用 numpy 二进制文件进行计算（没有理由重新发明已经润滑好的轮子:)。

【解决方案2】：

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

【讨论】：