3d阵列的Numpy elementwise产品

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