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

Question

I'm currently doing some memory-intensive text processing, for which I have to construct a sparse matrix of float32s with dimensions of ~ (2M, 5M) . I'm constructing this matrix column by column when reading a corpus of 5M documents. For this purpose I use a sparse dok_matrix data structure from SciPy . However, when arriving at the 500 000'th document, my memory is full (approx. 30GB is used) and the program crashes. What I eventually want to do, is perform a dimensionality reduction algorithm on the matrix using sklearn , but, as said, it is impossible to hold and construct the entire matrix in memory. I've looked into numpy.memmap , as sklearn supports this, and tried to memmap some of the underlying numpy data structures of the SciPy sparse matrix, but I could not succeed in doing this.

It is impossible for me to save the entire matrix in a dense format, since this would require 40TB of disk space. So I think that HDF5 and PyTables are no option for me (?).

My question is now: how can I construct a sparse matrix on the fly, but writing directly to disk instead of memory, and such that I can use it afterwards in sklearn?

Thanks!

Answer 1

We've come across similar problems in the field of single cell genomics data dealing with large sparse datasets on disk. I'll show you a small simple example of how I would deal with this. My assumptions are that you're very memory constrained, and probably can't fit multiple copies of the sparse matrix into memory at once. This will work even if you can't fit one entire copy.

I would construct an on disk sparse CSC matrix column by column. A sparse csc matrix uses 3 underlying arrays:

• data : the values stored in the matrix
• indices : the row index for each value in the matrix
• indptr : an array of length n_cols + 1 , which divides indices and data by which column they belong to.

As an explanatory example, the values for column i are stored in the range indptr[i]:indptr[i+1] of data . Similarly, the row indices for these values can be found by indices[indptr[i]:indptr[i+1]] .

To simulate your data generating process (parsing a document, I assume) I'll define a function process_document which returns the values for indices and data for the relevant document.

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()

Now I'll create a group in and hdf5 file which will store the constituent arrays of a sparse matrix.

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

And finally a function for reading this group as a sparse matrix (this one is pretty simple).

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"])

Now we'll create the on-disk sparse matrix and write one column at a time to it (I'm using fewer columns since this can be kinda slow).

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"])

Again this is considering a very memory constrained situation, where you might not be able to fit the entire sparse matrix in memory when creating it. A much faster way to do this, if you can handle the entire sparse matrix plus at least one copy, would be to not bother with the on disk storage (similar to other suggestions). However, using a slight modification of this code should give you better performance:

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))

This should be fairly fast, since it only makes a copy of the data once you concatenate the arrays. Other current posted solutions reallocated the arrays as you processed, making many copies of large arrays.

Answer 2

It would be great if you could provide a minimal working code. I can't see if your matrix gets too big by construction (1) or just because you have too much data (2). If you don't really care about building this matrix yourself, you can directly look at my remark 2.

For problem (1), in the example code below, I made a wrapper class to build a csr_matrix chunk by chunk. The idea is to just add (row,column,data) tuples of lists until a buffer limit (see remark 1) is reached, and actually update the matrix at this moment. When the limit is reached, it will reduce the data in memory since the csr_matrix constructor adds data that have the same (row,column) tuples. This part only allows you to construct the sparse matrix in a fast manner (much faster than creating a sparse matrix for each row) and avoids memory error due to the redundancy of the (row,column) when a word appears several times in a document.

import numpy as np
import scipy.sparse

class SparseMatrixBuilder():
    def __init__(self, shape, build_size_limit):
        self.sparse_matrix = scipy.sparse.csr_matrix(shape)
        self.shape = shape
        self.build_size_limit = build_size_limit
        self.data_temp = []
        self.col_indices_temp = []
        self.row_indices_temp = []


    def add(self, data, col_indices, row_indices):
        self.data_temp.append(data)
        self.col_indices_temp.append(col_indices)
        self.row_indices_temp.append(row_indices)
        if len(self.data_temp) == self.build_size_limit:
            self.sparse_matrix += scipy.sparse.csr_matrix(
                (np.concatenate(self.data_temp),
                 (np.concatenate(self.col_indices_temp),
                  np.concatenate(self.row_indices_temp))),
                shape=self.shape
            )
            self.data_temp = []
            self.col_indices_temp = []
            self.row_indices_temp = []

    def get_matrix(self):
        self.sparse_matrix += scipy.sparse.csr_matrix(
            (np.concatenate(self.data_temp),
             (np.concatenate(self.col_indices_temp),
              np.concatenate(self.row_indices_temp))),
            shape=self.shape
        )
        self.data_temp = []
        self.col_indices_temp = []
        self.row_indices_temp = []
        return self.sparse_matrix

For problem (2), you can easily extend this class by adding a save method that stores the matrix on disk once the limit (or a second limit) is reached. As such, you'll end up with multiple chunks of sparse matrices on disk. Then you'll need a dimensionality reduction algorithm that can handle chunked matrices (see remark 2).

remark 1: the buffer limit here is not really well defined. It would be better to check for the actual size of the numpy arrays data_temp, col_indices_temp and row_indices_temp compared to the RAM available on the machine (which is quite easy to automatize with python).

remark 2: gensim is a python library that has the great advantage to use chunked files for building NLP models. So you could build a dictionary, construct a sparse matrix and reduce it dimension with that library, without much RAM needed.

Answer 3

I'm assuming that all your data can fit in memory using a more memory-friendly sparse matrix format such as COO. If it does not, there is almost no hope you will be able to proceed with sklearn , even by using mmap . Indeed sklearn will likely create subsequent objects with memory requirements of the same order of magnitude as your input.

Scipy's dok_matrix are actually a sub-class of the vanilla dict . They store the data using individual python objects and tons of pointers, so they are not memory efficient. The most compact representation is the coo_matrix format. You can incrementally build the data required to create a COO matrix by pre-allocating arrays for the coordinates (rows and cols) and the data; and eventually increase these buffers if your initial guess was wrong.

def get_coo_from_iter(iterable, n_data_hint=1<<20, idx_dtype='uint32', data_dtype='float32'):
    counter = 0
    rows = numpy.empty(n_data_hint, dtype=idx_dtype)
    cols = numpy.empty(n_data_hint, dtype=idx_dtype)
    data = numpy.empty(n_data_hint, dtype=data_dtype)
    for row, col, value in iterable:
        if counter >= n_data_hint:
            n_data_hint *= 2
            rows, cols, data = _reallocate(rows, cols, data, n_data_hint)
        rows[counter] = row
        cols[counter] = col
        data[counter] = value
        counter += 1
    rows = rows[:counter]
    cols = cols[:counter]
    data = data[:counter]
    return coo_matrix((data, (rows, cols)))


def _reallocate(rows, cols, data, n):
    new_rows = numpy.empty(n, dtype=rows.dtype)
    new_cols = numpy.empty(n, dtype=cols.dtype)
    new_data = numpy.empty(n, dtype=data.dtype)
    new_rows[:rows.size] = rows
    new_cols[:cols.size] = cols
    new_data[:data.size] = data
    return new_rows, new_cols, new_data

which you can test with randomly-generated data like this:

def get_random_data(n, max_row=2000, max_col=5000):
    for _ in range(n):
        row = numpy.random.choice(max_row)
        col = numpy.random.choice(max_col)
        val = numpy.random.randn()
        yield row, col, val

# test when initial hint is good
coo = get_coo_from_iter(get_random_data(10000), n_data_hint=10000)
print(coo.shape)

# or to test when initial hint was too tiny
coo = get_coo_from_iter(get_random_data(10000), n_data_hint=1111)
print(coo.shape)

Once you have your COO matrix, you may want to convert to CSR using coo.tocsr() . The CSR matrices are more optimized for common operations such as dot product. It requires a bit more memory in the case where some rows were empty originally. This is because it stores pointers for all rows even empty ones.

Answer 4

看这里，最后他解释了如何将稀疏矩阵存储和直接读取到 Hdf5 文件。