python - 为什么矢量化版本更慢？

Question

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

谢谢

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 是一个优化的 C++ 函数，可以*利用您的 BLAS 实现。让我们看一下此操作的可比较时间:

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 实现可以进一步提高速度。

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