Numpy三维数组的逐元素相乘

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矩阵。

- Remy

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

- Divakar

---

谢谢Divakar。我希望我能使用这个，但是我被困在一个旧版本的numpy上，它没有einsum。是否有一个等效的tensordot调用？ - Remy

我从未能理解einsum如何工作。你是通过什么资料学习到语法的呢？ - rayryeng

@rayryeng 只是官方文档和进行多个小时的试验。还有许多其他事情，我看到别人在SO上使用它需要弄明白！就像你对permute所做的那样，玩一玩它。最近，通过玩弄，找到了如何使用np.einsum替换np.any。 - Divakar

@Remy，你最终有没有获得支持NumPy的np.einsum的访问权限？我不太确定如何在这种情况下让tensordot正常工作。 - Divakar

@Divakar 是的，我在自己的电脑上试过了，这真的快了很多，所以我会决定花点时间升级其他人的numpy。再次感谢。 - Remy

Answer 2

你只需要在张量的第一维上执行操作，该维由0标记：

c = tensordot(a, b, axes=(0,0))

这个方法可以按照你的期望工作。另外，你不需要一列轴，因为你只是沿着一个维度执行操作。使用axes([1,2],[2,1])来交叉乘法运算第二个和第三个维度。如果你用索引符号表示（爱因斯坦求和约定），这对应于c[i,j]=a[i,k,l]*b[j,k,l]，因此你正在收缩你想要保留的指数。

编辑：好的，问题在于两个三维对象的张量积是一个六维对象。由于缩并涉及成对的指数，因此通过tensordot操作不可能得到一个三维对象。诀窍在于将计算分为两部分：首先在要执行矩阵操作的索引上执行tensordot，然后取张量对角线以将你的四维对象减少到三维。在一条命令中：

d = np.diagonal(np.tensordot(a,b,axes=()), axis1=0, axis2=2)

在张量符号表示中，d[i,j,k] = c[i,j,i,k] = a[i,j,l]*b[i,l,k]。

- Mattia

---

我使用你的解决方案得到了一个(2, 2, 2, 2)的形状，而不是(10, 2, 2)，因此可能有些地方出错了。 - Holt

抱歉，我给了您一个错误的答案，请查看我的编辑。 - Mattia