3d 数组的 Numpy 元素乘积

Posted 2023-02-23

【中文标题】3d 数组的 Numpy 元素乘积

【英文标题】：Numpy elementwise product of 3d array

【发布时间】：2015-11-04 14:22:44

【问题描述】：

我有两个形状为 (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 矩阵。

【参考方案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

【讨论】：

谢谢迪瓦卡。我希望我可以使用它，但我有点坚持使用没有 einsum 的旧版本的 numpy。有 tensordot 等价调用吗？

我永远无法理解einsum 的工作原理。你读过哪些东西可以教你语法是如何工作的？

@rayryeng 只是官方文档和玩它的时间。我还看到人们用它做很多其他的事情，所以我需要弄清楚！就像你对permute 所做的那样，随便玩弄它。最近，又想通了how to use np.einsum to replace np.any。

@Remy 你终于可以访问np.einsum 支持的 NumPy 了吗？在这种情况下，我不确定如何让 tensordot 工作。

@Divakar 是的，我是在自己的计算机上完成的，这确实快得多，以至于我要咬紧牙关，升级其他人的 numpy。再次感谢。

【参考方案2】：

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

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

这将如您所愿。此外，您不需要轴列表，因为它只是沿着您正在执行操作的一个维度。使用axes([1,2],[2,1])，您可以交叉乘以第 2 维和第 3 维。如果你用索引符号（爱因斯坦求和约定）写它，这对应于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]。

【讨论】：

我的解决方案是(2, 2, 2, 2)，而不是(10, 2, 2)，所以可能有问题。

对不起，我给你一个错误的答案，请看我的编辑。