用 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 所说，它确实适用于正确的尺寸。我添加了一个我现在认为与我的问题更相关的问题。

Edit2：这是我试图做的一个非常简化的例子。我相信，问题在于定义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数组，或多或少 do 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)

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