python - 用很少的 RAM 优化我的大数据代码

Question

我保存了一个 120 GB 的文件(通过 pickle 以二进制格式保存)，其中包含大约 50,000 (600x600) 个二维 numpy 数组。我需要使用中位数堆叠所有这些数组。最简单的方法是将整个文件作为数组列表读取并使用 np.median(arrays, axis=0)。但是，我没有太多的 RAM 可以使用，所以这不是一个好的选择。

所以，我试着逐个像素地堆叠它们，就像我一次关注一个像素位置 (i, j) 一样，然后一个一个地读取每个数组，附加列表中给定位置的值。保存所有数组中某个位置的所有值后，我使用 np.median 然后只需将该值保存在列表中——最后将包含每个像素的中值位置。最后，我可以将其重新整形为 600x600，这样就完成了。代码如下。

import pickle
import time
import numpy as np

filename = 'images.dat' #contains my 50,000 2D numpy arrays

def stack_by_pixel(i, j):
    pixels_at_position = []
    with open(filename, 'rb') as f:
        while True:
            try:
                # Gather pixels at a given position
                array = pickle.load(f)
                pixels_at_position.append(array[i][j])
            except EOFError:
                break
    # Stacking at position (median)
    stacked_at_position = np.median(np.array(pixels_at_position))
    return stacked_at_position

# Form whole stacked image
stacked = []
for i in range(600):
    for j in range(600):
        t1 = time.time()
        stacked.append(stack_by_pixel(i, j))
        t2 = time.time()
        print('Done with element %d, %d: %f seconds' % (i, j, (t2-t1)))

stacked_image = np.reshape(stacked, (600,600))

在看到一些时间打印输出后，我意识到这是极其低效的。每次完成一个位置 (i, j) 大约需要 150 秒左右，这并不奇怪，因为它正在一个一个地读取大约 50,000 个数组。鉴于我的大型数组中有 360,000 个 (i, j) 位置，预计需要 22 个月才能完成!显然这是不可行的。但我有点不知所措，因为没有足够的 RAM 可用于读取整个文件。或者也许我可以一次保存数组的所有像素位置(每个位置一个单独的列表)，因为它一个一个地打开它们，但不会在 Python 中保存 360,000 个列表(大约 50,000 个元素长)使用很多还有 RAM？

欢迎就如何在不使用大量 RAM 的情况下显着加快运行速度提出任何建议。谢谢!

Answer 1

这是 numpy 的 memory mapped arrays 的完美用例. 内存映射数组允许您将磁盘上的 .npy 文件视为作为 numpy 数组加载到内存中，而不实际加载它。就这么简单

arr = np.load('filename', mmap_mode='r')

在大多数情况下，您可以将其视为任何其他数组。数组元素仅在需要时加载到内存中。不幸的是，一些快速实验表明 median 不能很好地处理内存映射数组*，它似乎仍然一次将大部分数据加载到内存中。所以 median(arr, 0) 可能不起作用。

但是，您仍然可以遍历每个索引并计算中位数而不会遇到内存问题。

[[np.median([arr[k][i][j] for k in range(50000)]) for i in range(600)] for j in range(600)]

其中 50,000 反射(reflect)了数组的总数。

没有为提取单个像素而对每个文件进行 unpickling 的开销，运行时间应该快得多(大约 360000 次)。

当然，这留下了创建包含所有数据的 .npy 文件的问题。可以创建一个文件，如下所示，

arr = np.lib.format.open_memmap(
    'filename',              # File to store in
    mode='w+',               # Specify to create the file and write to it
    dtype=float32,           # Change this to your data's type
    shape=(50000, 600, 600)  # Shape of resulting array
)

然后，像以前一样加载数据并将其存储到数组中(它只是在幕后将其写入磁盘)。

idx = 0
with open(filename, 'rb') as f:
    while True:
        try:
            arr[idx] = pickle.load(f)
            idx += 1
        except EOFError:
            break

让它运行几个小时，然后回到这个答案的开头，看看如何加载它并取中值。再简单不过了**。

*我刚刚在一个 7GB 的文件上进行了测试，取了 5,000,000 个元素的 1,500 个样本的中值，内存使用量约为 7GB，这表明整个数组可能已加载到内存中。不过，先尝试这种方式并没有坏处。如果其他人对内存映射数组的中位数有经验，请随时发表评论。

** 如果您相信互联网上的陌生人。