2

假设我有两个数组：

import numpy as np
a=np.array([[1,2],
            [3,4]])

b=np.array([[1,2],
            [3,4]])

我想按元素乘以数组然后对元素求和，即1*1 + 2*2 + 3*3 + 4*4 = 30我可以使用：

np.tensordot(a, b, axes=((-2,-1),(-2,-1)))
>>> array(30)

现在，假设数组a和b是 2×2×2 数组：

a=np.array([[[1, 2],
             [3, 4]],

            [[5, 6],
             [7, 8]]])

b=np.array([[[1, 2],
             [3, 4]],

            [[5, 6],
             [7, 8]]])

并且我想对每个组执行相同的操作，即[[1,2],[3,4]]与[[1,2],[3,4]]then 元素相加的时间，与[[5,6],[7,8]]. 结果应该是array([ 30, 174]), where30 = 1*1 + 2*2 + 3*3 + 4*4和174 = 5*5 + 6*6 + 7*7 + 8*8。有没有办法使用 tensordot 做到这一点？

PS
我明白在这种情况下你可以简单地使用 sum 或 einsum：

np.sum(a*b,axis=(-2,-1))
>>> array([ 30, 174])

np.einsum('ijk,ijk->i',a,b)
>>> array([ 30, 174])

但这只是一个简化的示例，我需要使用它，tensordot因为它更快。

谢谢你的帮助！！

4

1 回答 1

1

您可以使用：np.diag(np.tensordot(a, b, axes=((1, 2), (1, 2))))来获得您想要的结果。但是，np.tensordot在您的情况下，使用或矩阵乘法并不是一个好主意，因为它们所做的工作比需要的要多得多。它们被有效实现的事实并不能平衡它们进行的计算量比需要的多得多的事实（这里只有对角线有用）。np.einsum('ijk,ijk->i',a,b)在您的情况下，不会计算比需要更多的东西。你可以试试optimize=True甚至optimize='optimal'因为参数默认optimize设置为False。如果这还不够快，您可以尝试使用 NumExpr 来提高计算np.sum(a*b,axis=(1, 2))效率（可能是并行计算）。或者，您也可以使用 Numba 或 Cython。两者都支持快速并行循环。

于 2022-01-27T17:45:29.297 回答