编辑 下面我将保留我遇到的更加复杂的问题,但是我使用 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循环而不是使用繁琐的索引,可以大大加快速度:
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
(或其他替代方法)而不是传统的花式索引来保持速度优势?
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]
保存了img
的第i
层平面经过i
层LUT的第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
的以下不太美观的替代方案中任何一个都会运行得更快:
单个平面上的单个查找表运行速度约为x70倍:
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
当然,问题是如果没有使用任何Python循环,是否无法使用np.take
完成此操作?理想情况下,任何所需的重塑或调整大小都应该发生在LUT上,而不是图像上,但我可以接受你们能够想出来的任何方法...
lut = lut.reshape(planes, 256, 4)
吗?最后一个维度是4
。 - tzellekelut
实际上是一个 断点表,所以在我的代码中它被称为bkpt
,而我在翻译问题时漏掉了这个。 - Jaimenp.take
目前只有在两个输入都是 C 连续时才快(否则会复制它们)。你可能可以手动将 2D 数组转换为 1D 数组,但如果img
很大,这可能并不重要,而且如果值得折腾的话... - seberg