将大矩阵与 dask 相乘

Question

我正在做一个基本上归结为求解矩阵方程的项目

A.dot(x) = d

其中A 是一个矩阵，到 2000 年的尺寸大约为 10 000 000（我想最终在两个方向上增加它）。

A 显然不适合内存，所以必须并行化。我通过解决A.T.dot(A).dot(x) = A.T.dot(d) 来做到这一点。 A.T 的维度是 2000 乘以 2000。可以通过将 A 和 d 分成块 A_i 和 d_i，沿着行计算 A_i.T.dot(A_i) 和 A_i.T.dot(d_i)，然后将它们相加。非常适合并行化。我已经能够使用多处理模块来实现这一点，但是 1）由于内存使用，很难进一步扩展（在两个维度上增加 A），以及 2）不漂亮（因此不容易维护） .

Dask 似乎是一个非常有前途的库，可以解决这两个问题，我也做了一些尝试。我的A 矩阵计算起来很复杂：它基于大约 15 个不同的数组（大小等于A 中的行数），其中一些用于迭代算法来评估相关的勒让德函数。当chunk很小（10000行）时，构建任务图需要很长时间，并且占用大量内存（内存增加与迭代算法的调用不谋而合）。当块较大（50000 行）时，计算前的内存消耗会小很多，但在计算A.T.dot(A) 时会迅速耗尽。我试过cache.Chest，但它显着减慢了计算速度。

任务图一定很大很复杂——调用A._visualize()会崩溃。使用更简单的A 矩阵，它可以直接执行此操作（请参阅@MRocklin 的回复）。有没有办法让我简化它？

任何有关如何解决此问题的建议将不胜感激。

作为一个玩具示例，我尝试了

A = da.ones((2e3, 1e7), chunks = (2e3, 1e3)) 
(A.T.dot(A)).compute()

这也失败了，耗尽了所有内存，只有一个内核处于活动状态。使用chunks = (2e3, 1e5)，几乎所有内核都会立即启动，但MemoryError 会在 1 秒内出现（我目前的计算机上有 15 GB）。 chunks = (2e3, 1e4) 更有希望，但最终也消耗了所有内存。

编辑： 我删除了玩具示例测试，因为尺寸错误，并更正了其余的尺寸。正如@MRocklin 所说，它确实适用于正确的尺寸。我添加了一个我现在认为与我的问题更相关的问题。

编辑2： 这是我试图做的一个非常简化的例子。我相信，问题在于定义A 中的列所涉及的递归。

import dask.array as da

N = 1e6
M = 500

x = da.random.random((N, 1), chunks = 5*M)

# my actual 
A_dict = {0:x}
for i in range(1, M):
    A_dict[i] = 2*A_dict[i-1]
A = da.hstack(tuple(A_dict.values()))
A = A.rechunk((M*5, M))
ATA = A.T.dot(A)

这似乎导致了一个非常复杂的任务图，甚至在计算开始之前就占用了大量内存。

我现在已经解决了这个问题，方法是将递归放在一个函数中，使用numpy 数组，或多或少使用A = x.map_blocks(...)。

作为第二个说明，一旦我有了A 矩阵任务图，直接计算A.T.dot(A) 似乎确实会产生一些内存问题（内存使用情况不是很稳定）。因此，我明确地分块计算它，并对结果求和。即使有这些变通方法，dask 在速度和可读性方面也有很大的不同。

Answer 1

你的输出非常非常大。

>>> A.T.dot(A).shape
(10000000, 10000000)

也许您打算用另一个方向的转置来计算这个？

>>> A.dot(A.T).shape
(2000, 2000)

这仍然需要一段时间（这是一个大型计算），但它确实完成了。