3d 数组的 Numpy 元素乘积 [英] Numpy elementwise product of 3d array

问题描述

我有两个形状为 (N, 2, 2) 的 3d 数组 A 和 B，我想根据 N 轴按元素乘以每个 2x2 矩阵上的矩阵乘积.使用循环实现，它看起来像

I have two 3d arrays A and B with shape (N, 2, 2) that I would like to multiply element-wise according to the N-axis with a matrix product on each of the 2x2 matrix. With a loop implementation, it looks like

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

有没有办法不用循环就可以做到这一点?我已经研究了 tensordot，但一直无法让它工作.我想我可能想要类似 tensordot(a, b, axes=([1,2], [2,1])) 的东西，但这给了我一个 NxN 矩阵.

Is there a way I could do this without using a loop? I've looked into tensordot, but haven't been able to get it to work. I think I might want something like tensordot(a, b, axes=([1,2], [2,1])) but that's giving me an NxN matrix.

推荐答案

看来您正在对第一个轴上的每个切片进行矩阵乘法.同样，您可以使用 np.einsum 就像这样 -

It seems you are doing matrix-multiplications for each slice along the first axis. For the same, you can use np.einsum like so -

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

我们也可以使用np.matmul -

We can also use np.matmul -

np.matmul(A,B)

在 Python 3.x 上，这个 matmul 操作通过 @ 运算符 -

On Python 3.x, this matmul operation simplifies with @ operator -

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

这篇关于3d 数组的 Numpy 元素乘积的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！