Question

我有一个问题，我必须做以下计算。我想避免循环版本，所以我对它进行了矢量化。为什么循环版本实际上比矢量化版本快？有没有人对此有解释。

THX

import numpy as np
from numpy.core.umath_tests import inner1d

num_vertices = 40000
num_pca_dims = 1000
num_vert_coords = 3
a = np.arange(num_vert_coords * num_vertices * num_pca_dims).reshape((num_pca_dims, num_vertices*num_vert_coords)).T

#n-by-3
norms = np.arange(num_vertices * num_vert_coords).reshape(num_vertices,-1)

#Loop version
def slowversion(a,norms):
    res_list = []
    for c_idx in range(a.shape[1]):
        curr_col = a[:,c_idx].reshape(-1,3)
        res = inner1d(curr_col, norms)
        res_list.append(res)
    res_list_conc = np.column_stack(res_list)
    return res_list_conc


#Fast version
def fastversion(a,norms):
    a_3 = a.reshape(num_vertices, 3, num_pca_dims)
    fast_res = np.sum(a_3 * norms[:,:,None], axis=1)
    return fast_res


res_list_conc = slowversion(a,norms)
fast_res = fastversion(a,norms)
assert np.all(res_list_conc == fast_res)

Answer 1

您的“慢速代码”可能会做得更好，因为inner1d是一个可以*利用您的BLAS实现的单个优化C ++函数。让我们看看这个操作的可比时间：

np.allclose(inner1d(a[:,0].reshape(-1,3), norms), 
           np.sum(a[:,0].reshape(-1,3)*norms,axis=1))
True

%timeit inner1d(a[:,0].reshape(-1,3), norms)
10000 loops, best of 3: 200 µs per loop

%timeit np.sum(a[:,0].reshape(-1,3)*norms,axis=1)
1000 loops, best of 3: 625 µs per loop

%timeit np.einsum('ij,ij->i',a[:,0].reshape(-1,3), norms)
1000 loops, best of 3: 325 µs per loop

使用inner比纯粹的numpy操作快得多。请注意，einsum几乎是纯numpy表达式和good reason的两倍。由于您的循环不是很大并且大多数FLOPS都在inner计算中，因此inner操作的保存超过了循环的成本。

%timeit slowversion(a,norms)
1 loops, best of 3: 991 ms per loop

%timeit fastversion(a,norms)
1 loops, best of 3: 1.28 s per loop

#Thanks to DSM for writing this out
%timeit np.einsum('ijk,ij->ik',a.reshape(num_vertices, num_vert_coords, num_pca_dims), norms)
1 loops, best of 3: 488 ms per loop

把它重新组合在一起我们可以看到“慢速版”的整体优势胜出;但是，使用einsum实现，这种实现相当优化，可以进一步提高速度。

*我在代码中没有看到它，但它显然是线程化的。

为什么矢量化版本更慢？

1 个答案: