Numpy乘以2d矩阵的3d矩阵

Question

例如，我得到了形状（3,2,2）的矩阵A，例如

[
[[1,1],[1,1]], 
[[2,2],[2,2]], 
[[3,3],[3,3]]
]

形状（2,2）的矩阵B，例如，

[[1, 1], [0,1]]

我想实现形状c（3,2,2），如：

c = np.zeros((3,2,2))
for i in range(len(A)):
    c[i] = np.dot(B, A[i,:,:])

这使

[[[2. 2.]
  [1. 1.]]

 [[4. 4.]
 [2. 2.]]

 [[6. 6.]
 [3. 3.]]]

实现这一目标的最有效方法是什么？

谢谢。

Answer 1

使用np.tensordot然后交换轴。 所以，使用其中一个 -

np.tensordot(B,A,axes=((1),(1))).swapaxes(0,1)
np.tensordot(A,B,axes=((1),(1))).swapaxes(1,2)

我们可以在交换轴后重塑A到2D，使用带有np.dot二维矩阵乘法，重塑和交换轴可能会获得边际性能提升。

计时 -

# Original approach
def orgapp(A,B):
    m = A.shape[0]
    n = B.shape[0]
    r = A.shape[2]
    c = np.zeros((m,n,r))
    for i in range(len(A)):
        c[i] = np.dot(B, A[i,:,:])
    return c  

In [91]: n = 10000
    ...: A = np.random.rand(n,2,2)
    ...: B = np.random.rand(2,2)

In [92]: %timeit orgapp(A,B)
100 loops, best of 3: 12.2 ms per loop

In [93]: %timeit np.tensordot(B,A,axes=((1),(1))).swapaxes(0,1)
1000 loops, best of 3: 191 µs per loop

In [94]: %timeit np.tensordot(A,B,axes=((1),(1))).swapaxes(1,2)
1000 loops, best of 3: 208 µs per loop

# @Bitwise's solution
In [95]: %timeit np.flip(np.dot(A,B).transpose((0,2,1)),1)
1000 loops, best of 3: 697 µs per loop

Answer 2

另一种方案：

np.flip(np.dot(A,B).transpose((0,2,1)),1)