Why is vectorized version slower?

Question

I have a problem where I have to do the following calculation. I wanted to avoid the loop version, so I vectorized it. Why is the loop version actually fast than the vectorized version? Does anybody have an explanation for this.

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

Your "slow code" is likely doing better because inner1d is a single optimized C++ function that can* make use of your BLAS implementation. Lets look at comparable timings for this operation:

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

Using inner is quite a bit faster then the pure numpy operations. Note that einsum is almost twice as fast as pure numpy expressions and for good reason. As your loop is not that large and most of the FLOPS are in the inner computations the saving for the inner operation outweigh the cost of looping.

%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

Putting this back together we can see the overall advantage of the "slow version" wins out; however, using an einsum implementation, which is fairly optimized for this sort of thing, gives us a further speed increase.

*I don't see it right off in the code, but it is clearly threaded.