三维数组的Numpy elementwise产品。

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

我有两个3d数组A和B，它们的形状是(N, 2, 2)我想根据N轴按元素相乘，每个2x2矩阵上都有一个矩阵乘积。对于循环实现，它看起来是这样的

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.

有没有一种不用循环的方法?我研究过tensordot，但一直没能让它发挥作用。我想我可能需要一些像tensordot(a, b, axes=([1,2]，[2,1])这样的东西，但这给我一个NxN矩阵。

2 个解决方案

#1

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一样- - - - - -

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

We can also use np.matmul -

我们也可以用np。matmul -

np.matmul(A,B)

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

Python 3。这个matmul操作用@操作符-简化

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

#2

You just need to perform the operation on the first dimension of your tensors, which is labeled by 0:

你只需要在张量的第一个维度上执行这个操作，它被标记为0:

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

This will work as you wish. Also you don't need a list of axes, because it's just along one dimension you're performing the operation. With axes([1,2],[2,1]) you're cross multiplying the 2nd and 3rd dimensions. If you write it in index notation (Einstein summing convention) this corresponds to c[i,j] = a[i,k,l]*b[j,k,l], thus you're contracting the indices you want to keep.

这将如你所愿。你也不需要一个坐标轴的列表，因为它只是沿着一个维度执行操作。用坐标轴([1,2]，[2,1])交叉相乘。如果你用指数表示法(爱因斯坦求和约定)这对应于c[i,j] = a[i,k,l]*b[j,k,l]，因此你在收缩你想要保留的指标。

EDIT: Ok, the problem is that the tensor product of a two 3d object is a 6d object. Since contractions involve pairs of indices, there's no way you'll get a 3d object by a tensordot operation. The trick is to split your calculation in two: first you do the tensordot on the index to do the matrix operation and then you take a tensor diagonal in order to reduce your 4d object to 3d. In one command:

编辑:好的，问题是两个3d对象的张量积是一个6d对象。由于收缩涉及对指标，你不可能通过一个tensordot操作得到一个3d对象。诀窍是把你的计算分成两部分:首先你做索引上的tensordot来做矩阵运算，然后你取一个张量对角线来把你的4d对象减少到3d。在一个命令:

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

In tensor notation d[i,j,k] = c[i,j,i,k] = a[i,j,l]*b[i,l,k].

在张量符号d(i,j,k)= c(i,j,k)=(i,j l)* b(l,k)。