python - MemoryError on dask.linalg.solve(A, b) 求解 Ax = b

Question

我想用一个大的、对称的、密集的矩阵 A (>100GB) 求解方程 Ax=b。我正在尝试 Dask，因为它应该是 python 处理大数据的一个很好的高级 API。

这是我的代码:

n = 100000 # makes 75GB array
chunk_size = 5000
A = da.random.random(size=(n, n), chunks=(chunk_size, chunk_size))
b = np.random.random(n)
x = da.linalg.solve(A, b)
res = x.compute()

不幸的是，我收到以下错误，显然是因为 A 太大了。 Dask 中是否有任何智能机制可以防止 MemoryError(例如迭代解决问题)？

Traceback (most recent call last):
File "test_solver.py", line 27, in <module>
x = da.linalg.solve(A, b)
File "..\venv\lib\site-packages\dask\array\linalg.py", line 916, in solve
File "..\venv\lib\site-packages\dask\array\linalg.py", line 755, in lu
MemoryError

非常感谢您!

Answer 1

你的问题不是数组太大，而是你的 block 太小:数组由 1,000,000 block 组成，求解图包含的任务呈指数增长。

例如，如果您的 block 大小为 5000x5000(每个 block 190MB)，您会得到

In [5]: n = 100000 # makes 75GB array
   ...: chunk_size = 5000
   ...: A = da.random.random(size=(n, n), chunks=(chunk_size, chunk_size))
   ...: b = np.random.random(n)
   ...: x = da.linalg.solve(A, b)
In [7]: len(x.dask)
6541

图中有 6541 个任务。如果 block 大小为 1000x1000，您将获得 427401 个任务。

https://docs.dask.org/en/latest/best-practices.html