高效计算一维数组和三维数组沿一维的乘积 - NumPy

Question

我有两个numpy数组：

• 一个1D数组t，形状为(70L,)，元素称为ti
• 一个3D数组I，形状为(70L, 1024L, 1024L)，每个元素称为Ii。因此，Ii的维度为(1024L, 1024L)

我想在第一维上对这两个数组进行乘积，即：

tI = t1*I1,t2*I2,...,tN*IN

例如，可以再次获取一个新的维度为（70L，1024L，1024L）的数组，然后沿着第一维进行求和，以获得一个维度为（1024L，1024L）的数组。

tsum = t1*I1 + t2*I2 + ... +tN*IN

目前我对以下操作感到满意：

tI = np.asarray([t[i]*I[i,:,:] for i in range(t.shape[0])])
tsum = np.sum(tI,axis=0)

如果我的数组维度增加，它将变得有点慢。我想知道是否存在一种更优化于该特定任务的numpy或scipy函数？

非常感谢任何链接或信息。

格雷格

- gregory

Answer 1

您可以使用 np.tensordot函数 -

np.tensordot(t,I, axes=([0],[0]))

您可以使用 np.einsum 函数。

np.einsum('i,ijk->jk',t,I)

运行时测试和输出验证 -

In [21]: def original_app(t,I):
    ...:     tI = np.asarray([t[i]*I[i,:,:] for i in range(t.shape[0])])
    ...:     tsum = np.sum(tI,axis=0)
    ...:     return tsum
    ...: 

In [22]: # Inputs with random elements
    ...: t = np.random.rand(70,)
    ...: I = np.random.rand(70,1024,1024)
    ...: 

In [23]: np.allclose(original_app(t,I),np.tensordot(t,I, axes=([0],[0])))
Out[23]: True

In [24]: np.allclose(original_app(t,I),np.einsum('i,ijk->jk',t,I))
Out[24]: True

In [25]: %timeit np.tensordot(t,I, axes=([0],[0]))
1 loops, best of 3: 110 ms per loop

In [26]: %timeit np.einsum('i,ijk->jk',t,I)
1 loops, best of 3: 201 ms per loop

- Divakar

---

对我来说，“tensordot”实际上更快。 - egpbos

@egpbos 你试过使用相同的数据大小吗？我认为这可能与NumPy在安装过程中编译的方式有关。对于不同的硬件，它可能会表现出不同的性能，虽然这只是一个猜测。 - Divakar

是的，我只是复制了你的测试代码。对我来说，“tensordot”始终比“einsum”快约35％。'In [65]:％timeit np.einsum（'i，ijk-> jk'，t，I） 10个循环，3个中的最佳：每个循环106毫秒In [66]:％timeit np.tensordot（t，I，axes =（[0]，[0]）） 10个循环，3个中的最佳：每个循环65.9毫秒' - egpbos

@egpbos 嗯，有趣的结果。我猜硬件在其中扮演了重要角色。 - Divakar

嗨，谢谢。np.tensordot对我很有帮助，而且它比einsum更快。 - gregory

Answer 2

Divakar提供了最好（最有效）的答案。为了完整起见，另一种方法是使用Numpy的广播功能：

(t[:,np.newaxis,np.newaxis]*I).sum(axis=0)

通过在 t 中添加两个轴，可以进行广播操作，并且可以使用常规的 Numpy 操作，某些人可能会更易读。

- egpbos

---

而且在NumPy中，“广播”功能也是一个很棒的特性！喜欢那个。 - Divakar