NumPy:一次可用于许多小型矩阵的点积 [英] NumPy: Dot product for many small matrices at once

问题描述

我有一长串的3×3矩阵，例如，

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

import numpy as np

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

，对于每个小矩阵，我想执行一个外积dot(a, a.T).列表理解

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
    ])

可以，但是效果不佳.可能的改进可能是只制作大型dot产品，但是我在这里很难为其正确设置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.

有任何提示吗?

推荐答案

您可以将转置矩阵的列表获取为A.swapaxes(1, 2)，将所需的产品列表获取为A @ A.swapaxes(1, 2).

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

@运算符只是，有关文档表示"如果任一自变量为ND，N> 2，则将其视为驻留在最后两个索引中并据此广播的一组矩阵."

这篇关于NumPy:一次可用于许多小型矩阵的点积的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！