如何加速numpy.all和numpy.nonzero()函数？

Question

如何加速numpy.all和numpy.nonzero()函数？

6

我需要检查一个点是否在一个边界立方体内。需要检查的立方体数量非常大(约为4M)。我编写的代码如下:

import numpy as np

# set the numbers of points and cuboids
n_points = 64
n_cuboid = 4000000

# generate the test data
points = np.random.rand(1, 3, n_points)*512
cuboid_min = np.random.rand(n_cuboid, 3, 1)*512
cuboid_max = cuboid_min + np.random.rand(n_cuboid, 3, 1)*8

# main body: check if the points are inside the cuboids
inside_cuboid = np.all((points > cuboid_min) & (points < cuboid_max), axis=1)
indices = np.nonzero(inside_cuboid)

在我的电脑上，运行np.all需要8秒钟，而运行np.nonzero只需要3秒钟。有什么方法可以加速代码吗？

- f. c.

1个回答

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Divakar · Accepted Answer

我们可以通过沿着长度为3的最小轴切分来减少all-reduction 的内存拥塞，以获得inside_cuboid。

out = (points[0,0,:] > cuboid_min[:,0]) & (points[0,0,:] < cuboid_max[:,0]) & \
      (points[0,1,:] > cuboid_min[:,1]) & (points[0,1,:] < cuboid_max[:,1]) & \
      (points[0,2,:] > cuboid_min[:,2]) & (points[0,2,:] < cuboid_max[:,2])

时间 -

In [43]: %timeit np.all((points > cuboid_min) & (points < cuboid_max), axis=1)
2.49 s ± 20 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [51]: %%timeit
    ...: out = (points[0,0,:] > cuboid_min[:,0]) & (points[0,0,:] < cuboid_max[:,0]) & \
    ...:       (points[0,1,:] > cuboid_min[:,1]) & (points[0,1,:] < cuboid_max[:,1]) & \
    ...:       (points[0,2,:] > cuboid_min[:,2]) & (points[0,2,:] < cuboid_max[:,2])
1.95 s ± 10.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)