通过numpy加速多个矩阵产品

Question

在python中我有2个三维数组：

T，尺寸为(n,n,n)

U，尺寸为(k,n,n)

T和U可以看作是彼此相邻的许多二维数组。我需要将所有这些矩阵相乘，即我必须执行以下操作：

for i in range(n):
        H[:,:,i] =  U[:,:,i].dot(T[:,:,i]).dot(U[:,:,i].T)

由于n可能非常大，我想知道这个操作是否会以某种方式加速numpy。

Answer 1

仔细研究迭代器以及它们如何参与这些点积减少，我们可以将所有这些转换为一个np.einsum实现，如此 -

H = np.einsum('ijk,jlk,mlk->imk',U,T,U)