python - 在 Python 中即时在磁盘上构造稀疏矩阵

Question

我目前正在做一些内存密集型文本处理，为此我必须构造一个 float32s 的 sparse matrix，尺寸为 ~ (2M, 5M)。在阅读 5M 文档的语料库时，我正在逐列构建这个矩阵。为此，我使用了来自 SciPy 的稀疏 dok_matrix 数据结构。但是，当到达第 500 000 个文档时，我的内存已满(使用了大约 30GB)并且程序崩溃了。我最终想要做的是使用 sklearn 对矩阵执行降维算法，但是，如前所述，不可能在内存中保存和构造整个矩阵。我查看了 numpy.memmap，因为 sklearn 支持这一点，并尝试 memmap SciPy 稀疏矩阵的一些底层 numpy 数据结构，但我无法成功这样做。

我不可能以密集格式保存整个矩阵，因为这需要 40TB 的磁盘空间。所以我认为 HDF5 和 PyTables 不是我的选择(？)。

我现在的问题是:如何即时构建稀疏矩阵，但直接写入磁盘而不是内存，以便之后可以在 sklearn 中使用它？

谢谢!

Answer 1

我们在处理磁盘上的大型稀疏数据集的单细胞基因组数据领域遇到了类似的问题。我将向您展示一个简单的小例子，说明我将如何处理这个问题。我的假设是您的内存非常有限，并且可能无法一次将稀疏矩阵的多个副本放入内存中。即使您放不下一份完整的副本，这也可以使用。

我会逐列构造一个磁盘上的稀疏 CSC 矩阵。稀疏 csc 矩阵使用 3 个底层数组:

• data:存储在矩阵中的值
• indices:矩阵中每个值的行索引
• indptr:长度为n_cols + 1的数组，将indices和data除以哪一列属于。

作为一个说明性示例，列 i 的值存储在 data indptr[i]:indptr[i+1] 范围内。同样，这些值的行索引可以通过 indices[indptr[i]:indptr[i+1]] 找到。

为了模拟您的数据生成过程(我假设是解析文档)，我将定义一个函数 process_document，它返回 indices 和 data< 的值 用于相关文档。

import numpy as np
import h5py
from scipy import sparse

from tqdm import tqdm  # For monitoring the writing process
from typing import Tuple, Union  # Just for argument annotation

def process_document():
    """
    Simulate processing a document. Results in sparse vector represenation.
    """
    n_items = np.random.negative_binomial(2, .0001)
    indices = np.random.choice(2_000_000, n_items, replace=False)
    indices.sort()
    data = np.random.random(n_items).astype(np.float32)
    return indices, data

def data_generator(n):
    """Iterator which yields simulated data."""
    for i in range(n):
        yield process_document()

现在我将在 hdf5 文件中创建一个组，该文件将存储稀疏矩阵的组成数组。

def make_sparse_csc_group(f: Union[h5py.File, h5py.Group], groupname: str, shape: Tuple[int, int]):
    """
    Create a group in an hdf5 file that can store a CSC sparse matrix.
    """
    g = f.create_group(groupname)
    g.attrs["shape"] = shape
    g.create_dataset("indices", shape=(1,), dtype=np.int64, chunks=True, maxshape=(None,))
    g["indptr"] = np.zeros(shape[1] + 1, dtype=int) # We want this to have a zero for the first value
    g.create_dataset("data", shape=(1,), dtype=np.float32, chunks=True, maxshape=(None,))
    return g

最后是一个将这个组读取为稀疏矩阵的函数(这个非常简单)。

def read_sparse_csc_group(g: Union[h5py.File, h5py.Group]):
    return sparse.csc_matrix((g["data"], g["indices"], g["indptr"]), shape=g.attrs["shape"])

现在我们将创建磁盘上的稀疏矩阵并一次写入一列(我使用的列较少，因为这可能有点慢)。

N_COLS = 10

def make_disk_matrix(f, groupname, data_iter, shape):
    group = make_sparse_csc_group(f, "mtx", shape)

    indptr = group["indptr"]
    data = group["data"]
    indices = group["indices"]
    n_total = 0

    for doc_num, (cur_indices, cur_data) in enumerate(tqdm(data_iter)):
        n_cur = len(cur_indices)
        n_prev = n_total
        n_total += n_cur
        indices.resize((n_total,))
        data.resize((n_total,))
        indices[n_prev:] = cur_indices
        data[n_prev:] = cur_data
        indptr[doc_num+1] = n_total

# Writing
with h5py.File("data.h5", "w") as f:
    make_disk_matrix(f, "mtx", data_generator(10), (2_000_000, 10))

# Reading
with h5py.File("data.h5", "r") as f:
    mtx = read_sparse_csc_group(f["mtx"])

再次考虑到内存非常受限的情况，在这种情况下，您可能无法在创建时将整个稀疏矩阵放入内存中。如果您可以处理整个稀疏矩阵加上至少一个副本，那么执行此操作的一种更快的方法是不打扰磁盘存储(类似于其他建议)。但是，稍微修改一下这段代码应该会给你更好的性能:

def make_memory_mtx(data_iter, shape):
    indices_list = []
    data_list = []
    indptr = np.zeros(shape[1]+1, dtype=int)
    n_total = 0

    for doc_num, (cur_indices, cur_data) in enumerate(data_iter):
        n_cur = len(cur_indices)
        n_prev = n_total
        n_total += n_cur
        indices_list.append(cur_indices)
        data_list.append(cur_data)
        indptr[doc_num+1] = n_total

    indices = np.concatenate(indices_list)
    data = np.concatenate(data_list)

    return sparse.csc_matrix((data, indices, indptr), shape=shape)

mtx = make_memory_mtx(data_generator(10), shape=(2_000_000, 10))

这应该相当快，因为​​它只会在您连接数组后复制数据。当前发布的其他解决方案在您处理时重新分配了数组，从而制作了许多大型数组的副本。