使用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

但是,如果使用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(或替代方法)而不是传统的花式索引来保持速度优势？

---

原始问题 我有一组想要在图像上使用的查找表(LUT).保持LUT的阵列具有形状(planes, 256, n),并且图像具有形状(planes, rows, cols).两者都与LUT dtype = 'uint8'的256轴匹配.想法是p通过nLUT的p第-th平面的每个LUT 运行图像的第-th平面.

如果我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穿过jLUT第一i平面的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比跑得快得多:

1. 单个平面上的单个LUT快速运行x70:

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

2. 运行所有所需组合的python循环几乎快了x6:

   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循环没有办法做到这一点？理想情况下,无论需要重塑还是调整大小都应该在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,为第1面,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)