循环向量化 NumPy

Question

循环向量化 NumPy

5

我相对于Python还是新手，我有一个嵌套的for循环。由于for循环运行时间较长，我正在尝试找到一种向量化此代码的方法，以使其运行更快。

在这种情况下，coord是一个三维数组，其中coord[x, 0, 0]和coord[x, 0, 1]是整数，coord[x, 0, 2]是0或1。H是SciPy稀疏矩阵，x_dist、y_dist、z_dist和a都是浮点数。

# x_dist, y_dist, and z_dist are floats
# coord is a num x 1 x 3 numpy array where num can go into the hundreds of thousands
num = coord.shape[0]    
H = sparse.lil_matrix((num, num))
for i in xrange(num):
    for j in xrange(num):
        if (np.absolute(coord[i, 0, 0] - coord[j, 0, 0]) <= 2 and
                (np.absolute(coord[i, 0, 1] - coord[j, 0, 1]) <= 1)):

            x = ((coord[i, 0, 0] * x_dist + coord[i, 0, 2] * z_dist) -
                 (coord[j, 0, 0] * x_dist + coord[j, 0, 2] * z_dist))

            y = (coord[i, 0, 1] * y_dist) - (coord[j, 0, 1] * y_dist)

            if a - 0.5 <= np.sqrt(x ** 2 + y ** 2) <= a + 0.5:
                H[i, j] = -2.7

我还了解到使用NumPy进行广播虽然速度更快，但会导致大量的临时数组内存使用。是选择向量化路线更好，还是尝试使用类似Cython的东西？

- sonicxml

2个回答

2

计算的本质使得它难以使用我熟悉的numpy方法进行向量化。我认为在速度和内存使用方面最好的解决方案是cython。然而，您可以通过使用numba来获得一些加速。这里是一个示例（请注意，通常您使用autojit作为装饰器）：

import numpy as np
from scipy import sparse
import time
from numba.decorators import autojit
x_dist=.5
y_dist = .5
z_dist = .4
a = .6
coord = np.random.normal(size=(1000,1000,1000))

def run(coord, x_dist,y_dist, z_dist, a):
    num = coord.shape[0]    
    H = sparse.lil_matrix((num, num))
    for i in xrange(num):
        for j in xrange(num):
            if (np.absolute(coord[i, 0, 0] - coord[j, 0, 0]) <= 2 and
                    (np.absolute(coord[i, 0, 1] - coord[j, 0, 1]) <= 1)):

                x = ((coord[i, 0, 0] * x_dist + coord[i, 0, 2] * z_dist) -
                     (coord[j, 0, 0] * x_dist + coord[j, 0, 2] * z_dist))

                y = (coord[i, 0, 1] * y_dist) - (coord[j, 0, 1] * y_dist)

                if a - 0.5 <= np.sqrt(x ** 2 + y ** 2) <= a + 0.5:
                    H[i, j] = -2.7
    return H

runaj = autojit(run)

t0 = time.time()
run(coord,x_dist,y_dist, z_dist, a)
t1 = time.time()
print 'First Original Runtime:', t1 - t0

t0 = time.time()
run(coord,x_dist,y_dist, z_dist, a)
t1 = time.time()
print 'Second Original Runtime:', t1 - t0

t0 = time.time()
run(coord,x_dist,y_dist, z_dist, a)
t1 = time.time()
print 'Third Original Runtime:', t1 - t0

t0 = time.time()
runaj(coord,x_dist,y_dist, z_dist, a)
t1 = time.time()
print 'First Numba Runtime:', t1 - t0

t0 = time.time()
runaj(coord,x_dist,y_dist, z_dist, a)
t1 = time.time()
print 'Second Numba Runtime:', t1 - t0

t0 = time.time()
runaj(coord,x_dist,y_dist, z_dist, a)
t1 = time.time()
print 'Third Numba Runtime:', t1 - t0

我得到了这个输出：

First Original Runtime: 21.3574919701
Second Original Runtime: 15.7615520954
Third Original Runtime: 15.3634860516
First Numba Runtime: 9.87108802795
Second Numba Runtime: 9.32944011688
Third Numba Runtime: 9.32300305367

- jcrudy

谢谢你的建议！然而，当我尝试将其放入脚本中（使用@autojit装饰器）并使用IPython进行计时（%timeit %run Test.py），我得到的结果比普通Python慢。你有任何想法为什么会发生这种情况吗？ - sonicxml

@sonicxml 很有趣。你使用的是和我的示例相同的数据吗？Autojit需要为每种新类型的数据编译您传递给它的函数，并在运行时执行此操作。因此，对于小的示例，由于编译时间，它可能会更慢。这可能是您正在使用的示例的问题吗？ - jcrudy

啊，好的。我之前一直在用一个较小的数组进行测试，但是现在我将数组扩大后，Numba 比Python快多了。 - sonicxml

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

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

- Jaime · Accepted Answer

这是我对你的代码进行向量化的方式，稍后会讨论一些注意事项：

import numpy as np
import scipy.sparse as sps

idx = ((np.abs(coord[:, 0, 0] - coord[:, 0, 0, None]) <= 2) &
       (np.abs(coord[:, 0, 1] - coord[:, 0, 1, None]) <= 1))

rows, cols = np.nonzero(idx)
x = ((coord[rows, 0, 0]-coord[cols, 0, 0]) * x_dist +
     (coord[rows, 0, 2]-coord[cols, 0, 2]) * z_dist)
y = (coord[rows, 0, 1]-coord[cols, 0, 1]) * y_dist
r2 = x*x + y*y

idx = ((a - 0.5)**2 <= r2) & (r2 <= (a + 0.5)**2)

rows, cols = rows[idx], cols[idx]
data = np.repeat(2.7, len(rows))

H = sps.coo_matrix((data, (rows, cols)), shape=(num, num)).tolil()

正如您所指出的，问题将出现在第一个idx数组上，因为它的形状将为(num, num)，如果num达到“数十万”级别，很可能会导致内存不足。

一种潜在的解决方案是将问题分解成可管理的块。如果您有一个包含100,000个元素的数组，可以将其分成100个包含1,000个元素的块，并对每个10,000个块组合运行修改过的上述代码版本。您只需要一个1,000,000个元素的idx数组（您可以预先分配并重复使用以提高性能），并且您只需进行10,000次迭代循环，而不是当前实现的10,000,000,000次。这是一种简单的并行化方案，如果您有多核机器，则可以通过同时处理几个块来进一步改善。