使用Dask处理48万×48万协方差矩阵SVD时遇内存错误求助

我有一个以二进制格式存储在磁盘上的480,000×480,000大型协方差矩阵，需要计算对应的白化矩阵，因此需对该矩阵执行SVD（奇异值分解）。我的系统配置强劲（40核、360GB内存），但仍无法在内存中完成计算。尝试结合numpy.memmap与Dask来解决问题（Dask支持分块读取、处理数据并分块保存结果），但处理完整协方差矩阵时总是出现MemoryError，不清楚问题所在。

使用的简化代码 #

初始化Dask客户端与内存映射数组：

import numpy as np
import dask
import dask.array as da
from dask.distributed import LocalCluster, Client

client = Client(processes=False, n_workers=10, threads_per_worker=1, memory_limit="30GiB")

da_memmap = da.from_array(
    np.memmap("path_to_covariance_matrix.npy",
        mode="r",
        shape=(480000, 480000),
        dtype=np.float64
    ),
    chunks=(250, 480000)
)

执行SVD并存储结果：

u, s, v = da.linalg.svd(da_memmap)

u_path = "/mnt/big_storage/test_u.npy"
s_path = "/mnt/big_storage/test_s.npy"
v_path = "/mnt/big_storage/test_v.npy"

u_map = np.memmap(u_path, mode="w+", dtype=np.float64, shape=da_memmap.shape)
s_map = np.memmap(s_path, mode="w+", dtype=np.float64, shape=da_memmap.shape[0])
v_map = np.memmap(v_path, mode="w+", dtype=np.float64, shape=da_memmap.shape)

da.store([u, s, v], [u_map, s_map, v_map], lock=False, compute=True)

报错信息 #

2023-12-29 17:34:03,672 - distributed.protocol.pickle - ERROR - Failed to serialize <ToPickle: HighLevelGraph with 1 layers.
<dask.highlevelgraph.HighLevelGraph object at 0x786097501630>
 0. 132356258993664
>.
Traceback (most recent call last):
  File "/home/shash/miniconda3/envs/analysis/lib/python3.10/site-packages/distributed/protocol/pickle.py", line 63, in dumps
    result = pickle.dumps(x, **dump_kwargs)
MemoryError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/shash/miniconda3/envs/analysis/lib/python3.10/site-packages/distributed/protocol/pickle.py", line 68, in dumps
    pickler.dump(x)
MemoryError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/shash/miniconda3/envs/analysis/lib/python3.10/site-packages/distributed/protocol/pickle.py", line 81, in dumps
    result = cloudpickle.dumps(x, **dump_kwargs)
  File "/home/shash/miniconda3/envs/analysis/lib/python3.10/site-packages/cloudpickle/cloudpickle.py", line 1479, in dumps
    cp.dump(obj)
  File "/home/shash/miniconda3/envs/analysis/lib/python3.10/site-packages/cloudpickle/cloudpickle.py", line 1245, in dump
    return super().dump(obj)
MemoryError

使用可完全放入内存的较小协方差矩阵运行该代码可以正常工作，尝试调整chunk_size但无济于事。现在质疑Dask是否能够完成我想要的任务，相关的澄清也非常有用。

感谢帮助！

内容的提问来源于stack exchange，提问作者smashwhat