multiple numpy dot products without a loop

问题

Is it possible to compute several dot products without a loop? say you have the following:

a = randn(100, 3, 3)
b = randn(100, 3, 3)

I want to get an array z of shape (100, 3, 3) such that for all i

z[i, ...] == dot(a[i, ...], b[i, ...])

in other words, which verifies:

for va, vb, vz in izip(a, b, z):
    assert (vq == dot(va, vb)).all()

The straightforward solution would be:

z = array([dot(va, vb) for va, vb in zip(a, b)])

which uses an implicit loop (list comprehension + array).

Is there a more efficient way to compute z?

回答1:

np.einsum can be useful here. Try running this copy+pasteable code:

import numpy as np

a = np.random.randn(100, 3, 3)
b = np.random.randn(100, 3, 3)

z = np.einsum("ijk, ikl -> ijl", a, b)

z2 = np.array([ai.dot(bi) for ai, bi in zip(a, b)])

assert (z == z2).all()

einsum is compiled code and runs very fast, even compared to np.tensordot (which doesn't apply here exactly, but often is applicable). Here are some stats:

In [8]: %timeit z = np.einsum("ijk, ikl -> ijl", a, b)
10000 loops, best of 3: 105 us per loop


In [9]: %timeit z2 = np.array([ai.dot(bi) for ai, bi in zip(a, b)])
1000 loops, best of 3: 1.06 ms per loop

回答2:

Try Einstein summation in numpy:

z = np.einsum('...ij,...jk->...ik', a, b)

It's elegant and does not require you to write a loop, as you requested. It gives me a factor of 4.8 speed increase on my system:

%timeit z = array([dot(va, vb) for va, vb in zip(a, b)])
1000 loops, best of 3: 454 µs per loop

%timeit z = np.einsum('...ij,...jk->...ik', a, b)
10000 loops, best of 3: 94.6 µs per loop

回答3:

This solution still uses a loop, but is faster because it avoids unnecessary creation of temp arrays, by using the out arg of dot:

def dotloop(a,b):
    res = empty(a.shape)
    for ai,bi,resi in zip(a,b,res):
        np.dot(ai, bi, out = resi)
    return res

%timeit dotloop(a,b)
1000 loops, best of 3: 453 us per loop
%timeit array([dot(va, vb) for va, vb in zip(a, b)])
1000 loops, best of 3: 843 us per loop

回答4:

In addition to the other answers, I want to add that:

np.einsum("ijk, ijk -> ij", a, b)

Is suitable for a related case I encountered, where you have two 3D arrays consisting of matching 2D fields of 2D vectors (points or directions). This gives a kind of "element-wise" dot product between those 2D vectors.

For example:

np.einsum("ijk, ijk -> ij", [[[1,2],[3,4]]], [[[5,6],[7,8]]])
# => array([[17, 53]])

Where:

np.dot([1,2],[5,6])
# => 17
np.dot([3,4],[7,8])
# => 53

来源：https://stackoverflow.com/questions/24090889/multiple-numpy-dot-products-without-a-loop

标签

python

loops

numpy

matrix