3D数组的块状元素乘积

Question

我有两个形状为（N，2，2）的3d数组A和B，我想根据N轴在每个2x2矩阵上分别乘以矩阵乘积。通过循环实现，它看起来像

C[i] = dot(A[i], B[i])

有没有办法不用循环就可以做到这一点？我研究过Tensordot，但无法使其正常工作。我想我可能想要类似的东西，tensordot(a, b, axes=([1,2], [2,1]))但这给了我一个NxN矩阵。

Answer 1

似乎您正在沿第一个轴对每个切片进行矩阵乘法。同样，您可以这样使用np.einsum-

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

我们还可以使用np.matmul-

np.matmul(A,B)

在Python
3.x上，此matmul操作可简化为@运算符-

A @ B

标杆管理

方法-

def einsum_based(A,B):
    return np.einsum('ijk,ikl->ijl',A,B)

def matmul_based(A,B):
    return np.matmul(A,B)

def forloop(A,B):
    N = A.shape[0]
    C = np.zeros((N,2,2))
    for i in range(N):
        C[i] = np.dot(A[i], B[i])
    return C

时间-

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

In [45]: %timeit einsum_based(A,B)
    ...: %timeit matmul_based(A,B)
    ...: %timeit forloop(A,B)
100 loops, best of 3: 3.08 ms per loop
100 loops, best of 3: 3.04 ms per loop
100 loops, best of 3: 10.9 ms per loop