这两个向量相乘的最Pythonic方式是什么？

Question

我有两个形状为：

的ndarray：

A = (32,512,640)

B = (4,512)

我需要将A和B相乘，以便得到一个新的ndarray：

C = (4,32,512,640)

另一种思考方式是将向量B的每一行沿着A的axis=-2进行乘法运算，这会生成一个新的1,32,512,640的立方体。可以遍历B的每一行，生成1,32,512,640个立方体，然后使用np.concatenate或np.vstack来构建C，例如：

# Sample inputs, where the dimensions aren't necessarily known
a = np.arange(32*512*465, dtype='f4').reshape((32,512,465))
b = np.ones((4,512), dtype='f4')

# Using a loop
d = []
for row in b:
    d.append(np.expand_dims(row[None,:,None]*a, axis=0))

# Or using list comprehension
d = [np.expand_dims(row[None,:,None]*a,axis=0) for row in b]

# Stacking the final list
result = np.vstack(d)

但我想知道是否有可能使用类似于np.einsum或np.tensordot这样的方法来一次性进行矢量化。 我仍在学习如何使用这两种方法，所以我不确定它是否适用于此处。

谢谢！

- wolfblade87

---

你已经尝试了什么？它有效吗？如果没有，它以何种方式未能按预期执行？ - Engineero

已添加我已经尝试过的内容，尽管在下面发布了有帮助的答案。它的表现符合预期，但我怀疑可能有更好或更有效的方法来完成这个任务（无论是通过性能还是减少代码行数）。 - wolfblade87

顺便说一下，@wolfblade87，“A”，“B”和“C”不是向量；它们分别是3D、2D和4D数组。 - kmario23

Answer 1

我们可以通过使用 broadcasting 技术，并在 B 的维度上使用扩展操作None/np.newaxis来实现。

C = A * B[:,None,:,None]

使用 einsum，就可以这样做 -

C = np.einsum('ijk,lj->lijk',A,B)

这里没有进行求和约简，所以使用einsum不会比explicit-broadcasting更好。但我们正在寻找一种Pythonic解决方案，一旦我们度过它的字符串表示法，就可以使用它。

让我们进行一些计时来完成这件事。

In [15]: m,n,r,p = 32,512,640,4
    ...: A = np.random.rand(m,n,r)
    ...: B = np.random.rand(p,n)

In [16]: %timeit A * B[:,None,:,None]
10 loops, best of 3: 80.9 ms per loop

In [17]: %timeit np.einsum('ijk,lj->lijk',A,B)
10 loops, best of 3: 109 ms per loop

# Original soln
In [18]: %%timeit
    ...: d = []
    ...: for row in B:
    ...:     d.append(np.expand_dims(row[None,:,None]*A, axis=0))
    ...: 
    ...: result = np.vstack(d)
10 loops, best of 3: 130 ms per loop

利用多核技术

我们可以利用numexpr的多核能力，适用于算术运算和大数据处理，从而提高性能。让我们来测试一下它的速度 -

In [42]: import numexpr as ne

In [43]: B4D = B[:,None,:,None] # this is virtually free

In [44]: %timeit ne.evaluate('A*B4D')
10 loops, best of 3: 64.6 ms per loop

简单来说就是：ne.evaluate('A*B4D',{'A':A,'B4D' :B[:,None,:,None]})。

相关帖子，介绍如何控制多核功能。

- Divakar

---

为了让意图更加清晰，使用 np.newaxis 而不是 None。 - MaxNoe

谢谢你的回答！我不知道你可以这样扩展维度并进行广播！ - wolfblade87