用很少的内存优化我的大数据代码

时间:2018-08-30 17:02:14

标签: python arrays python-3.x numpy optimization

我保存了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个数组。考虑到我的大型阵列中有360,000个(i, j)职位,这预计需要22个月才能完成!显然这是不可行的。但是我有点无所适从,因为没有足够的RAM来读取整个文件。或者也许我可以一次保存所有像素位置(每个位置一个单独的列表)作为数组,因为它一个接一个地打开它们,但不会在Python中保存360,000个列表(长约50,000个元素),这会占用很多内存也一样?

欢迎提出任何建议,以使我可以在不使用大量RAM的情况下使运行速度大大加快。谢谢!

2 个答案:

答案 0 :(得分:2)

这是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个表示数组总数。

没有去掉每个文件仅提取单个像素的开销,运行时间应该更快(大约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的文件上对其进行了测试,对1500个样本的中位数为5,000,000个元素,内存使用量约为7GB,这表明整个阵列可能已加载到内存中。不过先尝试这种方式并没有什么坏处。如果其他任何人对混合数组的中值有经验,请随时发表评论。

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

答案 1 :(得分:1)

注意:我使用的是Python 2.x,将其移植到3.x并不难。


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

准备

为了对此进行测试,我编写了一个小脚本,生成一个类似于您的泡菜文件。我假设您输入的图像是灰度的,具有8位深度,并使用numpy.random.randint生成了10000个随机图像。

此脚本将作为基准,我们可以将其与预处理和处理阶段进行比较。

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 https://stackoverflow.com/questions/6774724/why-python-has-limit-for-count-of-file-handles
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。