Question

我有一个函数，我想将其余的循环向量化。我相信这是正确的，我对性能感到满意，但我想更多地了解矢量化代码。功能是：

def f(x, A, c):
    # A : d*d numpy array
    # c : length d numpy array
    # x : N x d or length d numpy array
    x = np.atleast_2d(x)
    b = np.zeros(x.shape[0], dtype=np.bool)
    for row in range(x.shape[0]):
        xmc = x[row, :] - c
        b[row] = xmc.dot(A).dot(xmc) <= 1
    return b

是否可以对函数进行矢量化并删除剩余的循环，同时保持相当简单？当循环中的独立计算无法很好地矢量化时，是否有任何指导原则？ N和d的典型值分别为10000和4。谢谢。

Answer 1

你可以像这样进行矢量化 -

xc = x-c
b_out = ((xc.dot(A))*xc).sum(1) <= 1

您还可以使用np.einsum -

xc = x-c
b_out = np.einsum('ij,jk,ik->i',xc,A,xc) <= 1

运行时测试 -

定义功能：

def org_app(x, A, c):
    x = np.atleast_2d(x)
    b = np.zeros(x.shape[0], dtype=np.bool)
    for row in range(x.shape[0]):
        xmc = x[row, :] - c
        b[row] = xmc.dot(A).dot(xmc) <= 1
    return b

def vectorized_app1(x,A,c):    
    xc = x-c
    return ((xc.dot(A))*xc).sum(1) <= 1

def vectorized_app2(x,A,c):    
    xc = x-c
    return np.einsum('ij,jk,ik->i',xc,A,xc) <= 1

时序：

In [266]: N = 20
     ...: d = 20
     ...: A = np.random.rand(d,d)
     ...: c = np.random.rand(d)
     ...: x = np.random.rand(N,d)
     ...: 

In [267]: %timeit org_app(x,A,c)
1000 loops, best of 3: 274 µs per loop

In [268]: %timeit vectorized_app1(x,A,c)
10000 loops, best of 3: 46 µs per loop

In [269]: %timeit vectorized_app2(x,A,c)
10000 loops, best of 3: 63.7 µs per loop

In [270]: N = 100
     ...: d = 100
     ...: A = np.random.rand(d,d)
     ...: c = np.random.rand(d)
     ...: x = np.random.rand(N,d)
     ...: 

In [271]: %timeit org_app(x,A,c)
100 loops, best of 3: 2.74 ms per loop

In [272]: %timeit vectorized_app1(x,A,c)
1000 loops, best of 3: 1.46 ms per loop

In [273]: %timeit vectorized_app2(x,A,c)
100 loops, best of 3: 4.72 ms per loop

点产品矢量化 - NumPy

1 个答案: