如何向量化这个python代码？

Question

我正在尝试使用NumPy和矢量化操作来使代码段运行得更快.然而,我似乎误解了如何对这段代码进行矢量化(可能是由于对矢量化的理解不完全).

这是带循环的工作代码(A和B是已设置大小的2D数组,已经初始化):

for k in range(num_v):
    B[:] = A[:]
    for i in range(num_v):
        for j in range(num_v):
            A[i][j] = min(B[i][j], B[i][k] + B[k][j])
return A

这是我尝试矢量化上面的代码:

for k in range(num_v):
    B = numpy.copy(A)
    A = numpy.minimum(B, B[:,k] + B[k,:])
return A

为了测试这些,我使用了以下代码,上面的代码包含在一个名为'algorithm'的函数中:

def setup_array(edges, num_v):
    r = range(1, num_v + 1)
    A = [[None for x in r] for y in r]  # or (numpy.ones((num_v, num_v)) * 1e10) for numpy
    for i in r:
        for j in r:
            val = 1e10
            if i == j:
                val = 0 
            elif (i,j) in edges:
                val = edges[(i,j)]
            A[i-1][j-1] = val 
    return A

A = setup_array({(1, 2): 2, (6, 4): 1, (3, 2): -3, (1, 3): 5, (3, 6): 5, (4, 5): 2, (3, 1): 4, (4, 3): 8, (3, 4): 6, (2, 4): -4, (6, 5): -5}, 6) 
B = []
algorithm(A, B, 6)

预期的结果,以及我得到的第一个代码是:

[[0, 2, 5, -2, 0, 10] 
 [8, 0, 4, -4, -2, 9]
 [4, -3, 0, -7, -5, 5]
 [12, 5, 8, 0, 2, 13]
 [10000000000.0, 9999999997.0, 10000000000.0, 9999999993.0, 0, 10000000000.0]
 [13, 6, 9, 1, -5, 0]]

第二个(矢量化)函数返回:

[[ 0. -4.  0.  0.  0.  0.]
 [ 0. -4.  0. -4.  0.  0.]
 [ 0. -4.  0.  0.  0.  0.]
 [ 0. -4.  0.  0.  0.  0.]
 [ 0. -4.  0.  0.  0.  0.]
 [ 0. -4.  0.  0. -5.  0.]]

我错过了什么？

Answer 1

通常,您希望对代码进行矢量化,因为您认为代码运行速度太慢.
如果您的代码太慢,那么我可以告诉您正确的索引将使它更快.
而不是A[i][j]你应该写A[i, j]- 这避免了(子)数组的瞬态副本.
由于您在代码的最内层循环中执行此操作,因此这可能非常昂贵.

看这里:

In [37]: timeit test[2][2]
1000000 loops, best of 3: 1.5 us per loop

In [38]: timeit test[2,2]
1000000 loops, best of 3: 639 ns per loop

在您的代码中始终如一地这样做 - 我坚信这已经解决了您的性能问题!

话说回来...

......这是我对如何进行矢量化的看法

for k in range(num_v):
    numpy.minimum(A, np.add.outer(A[:,k], A[k,:]), A)
return A

numpy.minimum将比较两个数组并返回元素两个元素中较小的一个.如果你传递第三个参数,它将获取输出.如果这是一个输入数组,则整个操作就位.

正如Peter de Rivay解释的那样,你的广播解决方案存在问题 - 但在数学上你想要做的是通过添加两个向量来获得某种外部产品.因此,您可以在add函数上使用外部操作.

NumPy的二进制ufuncs具有执行某些特殊向量化操作的特殊方法,如reduce,accumulate,sum和outer.

Answer 2

该问题是由行中的阵列广播引起的:

A = numpy.minimum(B, B[:,k] + B[k,:])

B的大小为6乘6,B [:,k]是具有6个元素的数组,B [k,:]是具有6个元素的数组.

(因为你使用的是numpy数组类型,B [:,k]和B [k,:]都返回一个形状为N的rank-1数组

Numpy自动更改大小以匹配:

1. 首先将B [:,k]添加到B [k,:]以产生具有6个元素的中间数组结果.(这不是你想要的)
2. 其次,通过重复行,将6个元素阵列广播为6乘6矩阵
3. 第三,计算原始矩阵和该广播矩阵的最小值.

这意味着你的numpy代码相当于:

for k in range(num_v):
   B[:] = A[:]
   C=[B[i][k]+B[k][i] for i in range(num_v)]
   for i in range(num_v):
      for j in range(num_v):
         A[i][j] = min(B[i][j], C[j])

修复代码的最简单方法是使用矩阵类型而不是数组类型:

A = numpy.matrix(A)
for k in range(num_v):
    A = numpy.minimum(A, A[:,k] + A[k,:])

矩阵类型使用更严格的广播规则,因此在这种情况下:

1. 通过重复列将[:,k]扩展到6乘6矩阵
2. 通过重复行,将[k,:]扩展为6乘6矩阵
3. 将广播的矩阵加在一起以形成6乘6的矩阵
4. 应用最小值