用少量RAM优化我的大数据代码

Question

我有一个120 GB的文件保存(二进制文件夹pickle),包含大约50,000(600x600)2d 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个阵列.鉴于(i, j)我的大型阵列中有360,000个位置,预计需要22个月才能完成!显然这是不可行的.但我有点不知所措,因为没有足够的RAM可供读取整个文件.或者也许我可以一次性保存所有像素位置(每个位置的单独列表),因为它逐个打开它们,但不会在Python中保存360,000个列表(大约50,000个元素长)使用了很多RAM也是？

欢迎提出任何建议,如果不使用大量RAM,我可以大大加快运行速度.谢谢!

Answer 1

注意：我使用 Python 2.x，将其移植到 3.x 应该不难。

我的想法很简单 - 磁盘空间充足，所以让我们进行一些预处理，并将大的 pickle 文件转换成更容易处理的小块。

准备

为了测试这一点，我编写了一个小脚本，生成一个与您的类似的 pickle 文件。我假设您的输入图像是灰度图像并且具有 8 位深度，并使用 生成 10000 个随机图像numpy.random.randint。

该脚本将充当基准，我们可以将其与预处理和处理阶段进行比较。

import numpy as np
import pickle
import time

IMAGE_WIDTH = 600
IMAGE_HEIGHT = 600
FILE_COUNT = 10000

t1 = time.time()

with open('data/raw_data.pickle', 'wb') as f:
    for i in range(FILE_COUNT):
        data = np.random.randint(256, size=IMAGE_WIDTH*IMAGE_HEIGHT, dtype=np.uint8)
        data = data.reshape(IMAGE_HEIGHT, IMAGE_WIDTH)
        pickle.dump(data, f)
        print i,

t2 = time.time()
print '\nDone in %0.3f seconds' % (t2 - t1)

在测试运行中，此脚本在 372 秒内完成，生成约 10 GB 的文件。

预处理

让我们逐行分割输入图像——我们将有 600 个文件，其中文件包含每个输入图像的N行。N我们可以使用二进制存储行数据numpy.ndarray.tofile（然后使用 加载这些文件numpy.fromfile）。

import numpy as np
import pickle
import time

# Increase open file limit
# See /sf/ask/474230711/
import win32file
win32file._setmaxstdio(1024)

IMAGE_WIDTH = 600
IMAGE_HEIGHT = 600
FILE_COUNT = 10000

t1 = time.time()

outfiles = []
for i in range(IMAGE_HEIGHT):
    outfilename = 'data/row_%03d.dat' % i
    outfiles.append(open(outfilename, 'wb'))


with open('data/raw_data.pickle', 'rb') as f:
    for i in range(FILE_COUNT):
        data = pickle.load(f)
        for j in range(IMAGE_HEIGHT):
            data[j].tofile(outfiles[j])
        print i,

for i in range(IMAGE_HEIGHT):
    outfiles[i].close()

t2 = time.time()
print '\nDone in %0.3f seconds' % (t2 - t1)

在测试运行中，该脚本在 134 秒内完成，生成 600 个文件，每个文件 600 万字节。它使用了约 30MB 或 RAM。

加工

很简单，只需使用 加载每个数组numpy.fromfile，然后使用numpy.median获取每列的中位数，将其缩减回单行，并将这些行累积在列表中。

最后，用于numpy.vstack重新组合中值图像。

import numpy as np
import time

IMAGE_WIDTH = 600
IMAGE_HEIGHT = 600

t1 = time.time()

result_rows = []

for i in range(IMAGE_HEIGHT):
    outfilename = 'data/row_%03d.dat' % i
    data = np.fromfile(outfilename, dtype=np.uint8).reshape(-1, IMAGE_WIDTH)
    median_row = np.median(data, axis=0)
    result_rows.append(median_row)
    print i,

result = np.vstack(result_rows)
print result

t2 = time.time()
print '\nDone in %0.3f seconds' % (t2 - t1)

在测试运行中，该脚本在 74 秒内完成。您甚至可以很容易地将其并行化，但这似乎不值得。该脚本使用了约 40MB 的 RAM。

考虑到这两个脚本都是线性的，所用的时间也应该线性缩放。对于 50000 张图像，预处理大约需要 11 分钟，最终处理大约需要 6 分钟。这是在 i7-4930K @ 3.4GHz 上，特意使用 32 位 Python。