3D数组的块状元素乘积

本文介绍了3D数组的块状元素乘积的处理方法，对大家解决问题具有一定的参考价值，需要的朋友们下面随着小编来一起学习吧！ 问题描述 我有两个形状为(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矩阵.

解决方案

似乎您正在沿第一个轴对每个切片进行矩阵乘法.同样，您可以使用 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

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])

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.

解决方案

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)

We can also use np.matmul -

np.matmul(A,B)

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

A @ B

Benchmarking

Approaches -

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

Timings -

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