python - 使用 numpy.take 进行更快的花式索引

Question

编辑 我保留了下面面临的更复杂的问题，但我的问题是 np.take可以更好地总结如下。假设你有一个数组 img形状(planes, rows) , 和另一个数组 lut形状(planes, 256) ，并且您想使用它们创建一个新数组 out形状(planes, rows) , 其中out[p,j] = lut[p, img[p, j]] .这可以通过如下花式索引来实现:

In [4]: %timeit lut[np.arange(planes).reshape(-1, 1), img]
1000 loops, best of 3: 471 us per loop

但是，如果您不使用花哨的索引，而是在 planes 上使用 take 和 python 循环事情可以大大加快:

In [6]: %timeit for _ in (lut[j].take(img[j]) for j in xrange(planes)) : pass
10000 loops, best of 3: 59 us per loop

可以lut和 img以某种方式重新排列，以便整个操作在没有 python 循环的情况下发生，但使用 numpy.take (或替代方法)而不是传统的花式索引来保持速度优势？

---

原始问题我有一组要在图像上使用的查找表 (LUT)。包含 LUT 的数组的形状为 (planes, 256, n) , 图像的形状为 (planes, rows, cols) .两者都是dtype = 'uint8' , 匹配256 LUT 的轴。这个想法是运行 p - 通过每个 n 的图像的第平面来自 p 的 LUT -LUT 的第 plane。

如果我的lut和 img以下是:

planes, rows, cols, n = 3, 4000, 4000, 4
lut = np.random.randint(-2**31, 2**31 - 1,
                        size=(planes * 256 * n // 4,)).view('uint8')
lut = lut.reshape(planes, 256, n)
img = np.random.randint(-2**31, 2**31 - 1,
                    size=(planes * rows * cols // 4,)).view('uint8')
img = img.reshape(planes, rows, cols)

在使用像这样的花式索引后，我可以实现我的目标

out = lut[np.arange(planes).reshape(-1, 1, 1), img]

这给了我一个形状数组 (planes, rows, cols, n) , 其中out[i, :, :, j]持有i -img的第一个平面贯穿 j -i 的 LUT -LUT 的第 plane...

一切都很好，除了这个:

In [2]: %timeit lut[np.arange(planes).reshape(-1, 1, 1), img]
1 loops, best of 3: 5.65 s per loop

这是完全 Not Acceptable ，特别是因为我有以下所有使用 np.take 的不太好看的替代品比跑得快得多:

1. 单个平面上的单个 LUT 运行速度大约快 70 倍:

   In [2]: %timeit np.take(lut[0, :, 0], img[0])
   10 loops, best of 3: 78.5 ms per loop
   

2. 运行所有所需组合的 python 循环完成速度几乎快了 6 倍:

   In [2]: %timeit for _ in (np.take(lut[j, :, k], img[j]) for j in xrange(planes) for k in xrange(n)) : pass
   1 loops, best of 3: 947 ms per loop
   

3. 甚至运行 LUT 和图像中的所有平面组合，然后丢弃 planes**2 - planes不需要的比花哨的索引更快:

   In [2]: %timeit np.take(lut, img, axis=1)[np.arange(planes), np.arange(planes)]
   1 loops, best of 3: 3.79 s per loop
   

4. 我能想出的最快组合有一个 python 循环遍历平面并更快地完成 x13:

   In [2]: %timeit for _ in (np.take(lut[j], img[j], axis=0) for j in xrange(planes)) : pass
   1 loops, best of 3: 434 ms per loop
   

当然，问题是是否没有办法用 np.take 做到这一点？没有任何 python 循环？理想情况下，任何需要的 reshape 或调整大小都应该发生在 LUT 上，而不是图像上，但我对你们能想出的任何事情持开放态度......

Answer 1

首先我要说的是我真的很喜欢你的问题。在不重新排列 LUT 或 IMG 的情况下，以下解决方案有效:

%timeit a=np.take(lut, img, axis=1)
# 1 loops, best of 3: 1.93s per loop

但是从结果中你必须查询对角线:a[0,0], a[1,1], a[2,2];得到你想要的。我试图找到一种方法来仅对对角元素进行索引，但仍然没有成功。

这里有一些方法可以重新排列您的 LUT 和 IMG:如果 IMG 中的索引为 0-255，第一个平面，第二个平面为 256-511，第三个平面为 512-767，则以下工作，但这会阻止您使用'uint8'，这可能是个大问题...:

lut2 = lut.reshape(-1,4)
%timeit np.take(lut2,img,axis=0)
# 1 loops, best of 3: 716 ms per loop
# or
%timeit np.take(lut2, img.flatten(), axis=0).reshape(3,4000,4000,4)
# 1 loops, best of 3: 709 ms per loop

在我的机器中，您的解决方案仍然是最佳选择，并且非常合适，因为您只需要对角线求值，即 plane1-plane1、plane2-plane2 和 plane3-plane3:

%timeit for _ in (np.take(lut[j], img[j], axis=0) for j in xrange(planes)) : pass
# 1 loops, best of 3: 677 ms per loop

我希望这可以让您了解更好的解决方案。最好使用 flatten() 寻找更多选项，以及与 np.apply_over_axes() 或 np.apply_along_axis() 类似的方法，这似乎很有希望。

我使用下面的代码生成数据:

import numpy as np
num = 4000
planes, rows, cols, n = 3, num, num, 4
lut = np.random.randint(-2**31, 2**31-1,size=(planes*256*n//4,)).view('uint8')
lut = lut.reshape(planes, 256, n)
img = np.random.randint(-2**31, 2**31-1,size=(planes*rows*cols//4,)).view('uint8')
img = img.reshape(planes, rows, cols)