numpy memmap 内存使用 - 想要迭代一次

Question

假设我在磁盘上保存了一些大矩阵。将它全部存储在内存中是不可行的，所以我使用 memmap 来访问它

A = np.memmap(filename, dtype='float32', mode='r', shape=(3000000,162))

现在假设我想迭代这个矩阵（本质上不是以有序的方式），这样每一行将被访问一次。

p = some_permutation_of_0_to_2999999()

我想做这样的事情：

start = 0
end = 3000000
num_rows_to_load_at_once = some_size_that_will_fit_in_memory()
while start < end:
    indices_to_access = p[start:start+num_rows_to_load_at_once]
    do_stuff_with(A[indices_to_access, :])
    start = min(end, start+num_rows_to_load_at_once)

随着这个过程的进行，我的计算机变得越来越慢，我的 RAM 和虚拟内存使用量呈爆炸式增长。

有没有办法强制 np.memmap 使用一定数量的内存？ （我知道我需要的行数不会超过我计划一次读取的行数，而且缓存不会真正帮助我，因为我只访问每一行一次）

也许有其他方法可以按自定义顺序在 np 数组上进行迭代（类似生成器）？我可以使用 file.seek 手动编写它，但它恰好比 np.memmap 实现慢得多

do_stuff_with() 不保留对其接收到的数组的任何引用，因此在这方面没有“内存泄漏”

谢谢

Answer 1

这是我一段时间以来一直试图解决的问题。我使用大型图像数据集，numpy.memmap 为使用这些大型数据集提供了一个方便的解决方案。

但是，正如您所指出的，如果我需要访问每个帧（或在您的情况下为行）以执行某些操作，RAM 使用量最终会达到最大值。

幸运的是，我最近找到了一个解决方案，它允许您在限制 RAM 使用量的同时遍历整个 memmap 数组。

解决方案：

import numpy as np

# create a memmap array
input = np.memmap('input', dtype='uint16', shape=(10000,800,800), mode='w+')

# create a memmap array to store the output
output = np.memmap('output', dtype='uint16', shape=(10000,800,800), mode='w+')

def iterate_efficiently(input, output, chunk_size):
    # create an empty array to hold each chunk
    # the size of this array will determine the amount of RAM usage
    holder = np.zeros([chunk_size,800,800], dtype='uint16')

    # iterate through the input, replace with ones, and write to output
    for i in range(input.shape[0]):
        if i % chunk_size == 0:
            holder[:] = input[i:i+chunk_size] # read in chunk from input
            holder += 5 # perform some operation
            output[i:i+chunk_size] = holder # write chunk to output

def iterate_inefficiently(input, output):
    output[:] = input[:] + 5

计时结果：

In [11]: %timeit iterate_efficiently(input,output,1000)
1 loop, best of 3: 1min 48s per loop

In [12]: %timeit iterate_inefficiently(input,output)
1 loop, best of 3: 2min 22s per loop

磁盘上的阵列大小约为 12GB。使用iterate_efficiently 函数将内存使用量保持在 1.28GB，而iterate_inefficiently 函数最终在 RAM 中达到 12GB。

这是在 Mac OS 上测试的。

Answer 2

我已经尝试这个问题几天了，似乎有两种方法可以使用np.mmap 来控制内存消耗。第一个是可靠的，而第二个需要一些测试并且取决于操作系统。

选项 1 - 在每次读/写时重建内存映射：

def MoveMMapNPArray(data, output_filename):
    CHUNK_SIZE = 4096
    for idx in range(0,x.shape[1],CHUNK_SIZE):
        x = np.memmap(data.filename, dtype=data.dtype, mode='r', shape=data.shape, order='F')
        y = np.memmap(output_filename, dtype=data.dtype, mode='r+', shape=data.shape, order='F')
        end = min(idx+CHUNK_SIZE, data.shape[1])
        y[:,idx:end] = x[:,idx:end]

其中data 的类型为np.memmap。每次读取都会丢弃 memmap 对象，这样可以防止数组被收集到内存中，并且如果块大小很小，则内存消耗会非常低。它可能会引入一些 CPU 开销，但在我的设置 (MacOS) 上发现它很小。

选项 2 - 自己构建 mmap 缓冲区并提供内存建议

如果你查看np.memmap 源代码here，你可以看到创建自己的memmapped numpy 数组相对容易。具体来说，使用 sn-p：

mm = mmap.mmap(fid.fileno(), bytes, access=acc, offset=start)
mmap_np_array = ndarray.__new__(subtype, shape, dtype=descr, buffer=mm, offset=array_offset, order=order)

请注意，此 python mmap 实例存储为 np.memmap 的私有 _mmap 属性。

通过访问 python mmap 对象和 python 3.8，您可以使用其madvise 方法，描述为here。

这允许您建议操作系统在可用的情况下释放内存。 here for linux 描述了各种 madvise 常量，并指定了一些通用的跨平台选项。

MADV_DONTDUMP 常量看起来很有希望，但我还没有像选项 1 那样测试内存消耗。