Question

我有两个具有形状（N，2,2）的3d阵列A和B，我想根据N轴在每个2x2矩阵上使用矩阵乘积。使用循环实现，它看起来像

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

有没有办法可以在不使用循环的情况下执行此操作？我已经看过了腾讯，但是还没有能够让它发挥作用。我想我可能想要像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操作使用@ 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

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]，因此你要收缩你想要保留的指数。

编辑：好的，问题是两个3d对象的张量积是一个6d对象。由于收缩涉及成对索引，因此您无法通过tensordot操作获得3d对象。诀窍是将计算分成两部分：首先对索引执行tensordot以执行矩阵运算，然后采用张量对角线以将4d对象减少为3d。在一个命令中：

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

3d阵列的Numpy elementwise产品

2 个答案:

基准