NumPy: Dot product for many small matrices at once

问题

I have a long array of 3-by-3 matrices, e.g.,

import numpy as np

A = np.random.rand(25, 3, 3)

and for each of the small matrices, I would like to perform an outer product dot(a, a.T). The list comprehension

import numpy as np

B = np.array([
    np.dot(a, a.T) for a in A
    ])

works, but doesn't perform well. A possible improvement could be to do just one big dot product, but I'm having troubles here setting up A correctly for it.

Any hints?

回答1:

You can obtain the list of transposed matrices as A.swapaxes(1, 2), and the list of products you want as A @ A.swapaxes(1, 2).

import numpy as np

A = np.random.rand(25, 3, 3)

B = np.array([
    np.dot(a, a.T) for a in A
    ])

C = A @ A.swapaxes(1, 2)

(B==C).all()     # => True

The @ operator is just syntactic sugar for np.matmul, about which the documentation says that "if either argument is N-D, N > 2, it is treated as a stack of matrices residing in the last two indexes and broadcast accordingly."

回答2:

Looking closely at the code and trying to get a bigger picture reveals some information there. The first axis needs alignment between A and its transposed version as we are iterating through the first axis for both these versions in the code. Then, there is reduction along the third axis in both A and its transposed version. Now, to maintain the alignment and perform such a dot product reduction, there's one good vectorized option in np.einsum. Thus, to solve our case, the implementation would look like this -

B = np.einsum('ijk,ilk->ijl',A,A)

来源：https://stackoverflow.com/questions/38110296/numpy-dot-product-for-many-small-matrices-at-once

标签

python

arrays

numpy

vectorization

dot-product