如何以某种并行的方式有效地计算两个数组列表的外积？

Posted 2021-04-05

我有两个向量列表。

A = np.random.rand(100,2000)
B = np.random.rand(100,1000)

我需要用B的第一个条目计算A的第一个条目的外积。然后是第二个，然后是第三个，依此类推。

一个天真的循环

outers = []
for a, b in zip(A,B):
    outers.append(np.outer(a,b))

在我的电脑上需要≈730[ms]（通过&&timeit）。

最后，outers是一个100条长的2000x1000阵列，这是正确的。

必须有一种更有效的方法来并行化这项任务，因为现在我们实际上首先使用A[0]和THEN B[0]来计算A[1] B[1]，我们可以独立地并行地完成它们。

答案

如果你想并行执行NumPy数组操作，Dask是一个很好的选择。例如，您可以执行以下操作：

import dask.array as da
dA = da.from_array(A, chunks=(10, A.shape[1]))
dB = da.from_array(B, chunks=(10, B.shape[1]))

task_graph = dA[:, :, None] * dB[:, None]

result = task_graph.compute()

compute()步骤足够灵活，可以在多个线程，多个进程，多个核心，多个机器等上应用计算。

对于你问题中的特定示例，你不会在串行方法上获得太多收益，因为与简单地做100个外部产品的成本相比，分块输入数组和连接输出数组所涉及的开销很大。但是，对于更大的问题，这种方法可以带来显着的加速。