使用numpy.take进行更快的花式索引
编辑 我下面保留了我面临的更复杂的问题，但我对 np.take 的问题可以更简单地总结如下。假设你有一个形状为 (planes, rows) 的数组 img，还有一个形状为 (planes, 256) 的数组 lut，你想用它们创建一个形状为 (planes, rows) 的新数组 out，其中 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
但是，如果你不使用花式索引，而是用 take 和一个 Python 循环来处理 planes，那么速度会大大提升：
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（或其他方法）来保持速度优势，而不是传统的花式索引？
原始问题
我有一组查找表（LUTs），想用在一张图像上。存放 LUTs 的数组形状为 (planes, 256, n)，而图像的形状为 (planes, rows, cols)。这两个数组的类型都是 dtype = 'uint8'，与 LUT 的 256 轴相匹配。我的想法是将图像的 p 平面通过 LUT 的 p 平面中的每个 n 个 LUT 进行处理。
如果我的 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 个平面通过第 j 个 LUT 处理后的结果...
一切都很好，除了这一点：
In [2]: %timeit lut[np.arange(planes).reshape(-1, 1, 1), img]
1 loops, best of 3: 5.65 s per loop
这完全不可接受，尤其是因为我有以下几种看起来不太好的替代方案，使用 np.take 运行得更快：
在单个平面上使用单个 LUT 的速度快了大约 70 倍：
In [2]: %timeit np.take(lut[0, :, 0], img[0]) 10 loops, best of 3: 78.5 ms per loop一个 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甚至运行 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我能想到的最快组合是一个 Python 循环遍历平面，速度快了 13 倍：
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 循环来做到这一点？理想情况下，任何需要的重塑或调整大小应该在 LUT 上进行，而不是在图像上，但我对你们能想到的任何方法持开放态度...
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
在我的机器上，你的解决方案仍然是最佳选择，而且非常合适，因为你只需要对角线的评估，也就是平面1与平面1、平面2与平面2、平面3与平面3的比较：
%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)