如何有效地将2d numpy数组保存到磁盘

Question

我有一个巨大的2d numpy数组，应该作为共现矩阵工作。我试过用scipy。像我的数据结构一样稀疏，但是dok_matrix索引非常慢(慢4倍)。

# Impossible
import numpy
N = 1000000 (1 milion)
coo = np.zeros((N, N), dtype=np.uint32)

我想持久化这个数组。

在寻找保存它的方法后，我尝试使用PyTables或hd5py，但我找不到一种方法来保存它而不会耗尽内存。

with open(name, 'w') as _file:
   np.save(_file, coo)

例如，使用PyTables:

    import tables
    _file = tables.openFile(
                name,
                mode='w',
                title='Co-occurrence matrix')
    atom = tables.Atom.from_dtype(coo.dtype)
    _filters = tables.Filters(complib='blosc', complevel=5)
    ds = _file.createEArray(
            _file.root,
            'coo_matrix',
            atom,
            shape=(0, coo.shape[-1]),
            expectedrows=coo.shape[-1],
            filters=_filters)
    # ds[:] = coo => not an option
    for _index, _data in enumerate(coo):
        ds.append(coo[_index][np.newaxis,:])
    _file.close()

和使用hd5py:

import h5py
h5f = h5py.File(name, 'w')
h5f.create_dataset('dataset_1', data=coo)

这两种方法都在增加内存使用，直到我不得不终止进程。那么，有没有什么办法可以循序渐进呢?如果不可能，你能推荐另一种方法来保存这个矩阵吗?

编辑

我要创建这样的共现矩阵:

    coo = np.zeros((N, N), dtype=np.uint32)
    for doc_id, doc in enumerate(self.w.get_docs()):
        for w1, w2 in combinations(doc, 2):
                if w1 != w2:
                    coo[w1, w2] += 1

我想保存coo (2d numpy数组)，以便稍后从磁盘检索它并找到共现值，如:coo[w1, w2]

Answer 1

np.save是一种快速、有效的保存密集数组的方法。它所做的只是写一个小头，然后写入数组的数据缓冲区。

但是对于一个大数组，该数据缓冲区将具有N*N*4(对于您的dtype)字节-在一个连续的内存块中。这种设计也有利于元素访问——代码知道i,j元素的确切位置。

注意np.zeros((N,N))不会一次分配所有必需的内存。内存使用可能在使用过程中增长(包括保存)

np.savez对数据存储没有帮助。它对每个变量执行save，并将结果文件收集到zip存档(也可以压缩)中。

Tables和h5py可以保存和加载数据块，但是如果在创建或使用的时候需要在内存中存储整个数组，那么这就没有帮助了。

由于您的数组将非常稀疏，scipy稀疏矩阵可以节省内存，因为它只存储非零元素。但是它也要存储那个元素的坐标，所以每个非零元素的存储不是那么紧凑。有许多格式，每种格式都有其优点和缺点。

dok使用Python字典来存储数据，键值形式为(i,j)。它是增量构建稀疏矩阵的较好格式之一。我在其他SO问题中发现，使用dok的元素访问比使用普通字典慢。构建一个普通的字典会更快，然后update和dok。

lil是增量构建的另一种好格式。它将数据存储在2个列表的列表中。

一旦有了完整的i,j,data数组集，

coo对于构建矩阵很方便。

csr和csc适用于计算(特别是线性代数类)和元素访问。但是对于改变稀疏性(添加非零元素)没有好处。

但是你可以用一种格式建立一个矩阵，并很容易地将其转换为另一种格式来使用或存储。

关于存储稀疏矩阵有很多问题。最简单的是使用MATLAB兼容的。mat格式(csc用于稀疏)。要使用np.save，您需要保存底层数组(对于coo, csc, csr格式)。Python pickle必须用于保存dok或lil。

在[scipy] large sparse上搜索一下，看看关于这种矩阵的其他SO问题。您不是第一个使用numpy/scipy进行文档共现计算的人(它是scipy稀疏的3种主要用途之一，其他两种是线性代数和机器学习)。