python - 如何(/if)使用 dask 转置分布式 3D numpy 数组?

标签 python numpy mpi dask

我的问题是在 3D numpy 数组 A 上执行 3 次矩阵乘法,该数组太大而无法容纳在单个处理器中。在张量形式中,我想要 A_ijk B_km C_jn D_ip (B、C 和 D 都可以放入内存中)。我想知道 dask 是否适合此任务(或者其他工具是否更适合)。

我认为最好的方法是将此操作拆分为每个乘法,并确保它们都是本地的。此链接有一个非常有用的图表,总结了我正在谈论的内容 http://www.2decomp.org/1d_mode.html .

更详细地说:首先,要执行 A_ijk B_km,我应该将 A 分布在前两个轴上,并在本地对每支铅笔执行矩阵乘法(图中的第一步)。

然后,我需要转置数组,使 j 轴位于每个处理器的本地(并在 k(现在是 m)轴上分割),然后执行下一个乘法。 (因此从图中的第一步到第二步)。这就是我想知道 dask 是否可以提供帮助的地方。

我知道原则上可以使用 mpi4py 来完成此操作,但步骤非常重要,而 dask 数组具有有用的重新分块和转置方法,感觉与此应用程序相关。

这看起来是不是很适合 dask?

如果没有,是否有人知道任何可以执行这些步骤的 python 库?我知道fftw有执行此操作的例程,但我不知道如何编写必要的 C 代码,或者如何让它与 python 和 numpy 交互。

感谢您的帮助。

最佳答案

对于 future 的其他人来说,mpi4py 确实有一个转置方法。但它被称为 Alltoall/Alltoallv。 mpi4py 的文档或教程中没有对此进行解释。我在另一个教程中发现了它:https://info.gwdg.de/wiki/doku.php?id=wiki:hpc:mpi4py .

关于python - 如何(/if)使用 dask 转置分布式 3D numpy 数组?,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/56736851/

相关文章:

python - 使用 pandas python 从文件中读取特定日期行

python - 使用 matplotlib 中的数据框绘制 3D 图

performance - 可分配数组性能

c++ - MPI Allgatherv 函数出现问题

python - 如何从图像中删除小的独立蒙版?

python - 在第三次出现字符 python 后剥离字符串

python - 高效切片三角稀疏矩阵

python - 有一个 numpy.dot 文档示例不清楚,具有高维点积

python, numpy bool 数组 : negation in where statement

io - 使用 MPI IO 并行输出到单个文件