Python中对大型numpy数组进行内存高效排序

18

我需要使用numpy对一个非常大的基因组数据集进行排序。我有一个由26亿个浮点数组成的数组,维度为(868940742,3),一旦加载并静止,它就占用我的机器上约20GB的内存。我有一台带有16GB RAM,500GB固态硬盘和3.1GHz英特尔i7处理器的2015年早期13寸MacBook Pro。仅仅加载这个数组就会溢出到虚拟内存,但还没有到让我的机器受影响或不得不停止其他正在进行的工作的地步。

我逐步从22个较小的(N,2)子数组构建了这个非常大的数组。

函数FUN_1使用每个名为sub_arr的22个子数组生成2个新的(N,1)数组。

FUN_1的第一个输出是通过在b = array([X, F(X)])数组上插值sub_arr [:,0]中的值生成的,第二个输出是通过使用r = array([X, BIN(X)])数组将sub_arr [:,0]放入箱子中生成的。我分别称这些输出为b_arrrate_arr。该函数返回一个由(N,1)数组组成的3元组:

import numpy as np

def FUN_1(sub_arr):
    """interpolate b values and rates based on position in sub_arr"""

    b = np.load(bfile)
    r = np.load(rfile)

    b_arr = np.interp(sub_arr[:,0], b[:,0], b[:,1])
    rate_arr = np.searchsorted(r[:,0], sub_arr[:,0])  # HUGE efficiency gain over np.digitize...

    return r[rate_r, 1], b_arr, sub_arr[:,1] 

在for循环中,我调用这个函数22次,并使用值填充一个预先分配的全0数组full_arr = numpy.zeros([868940742, 3])

full_arr[:,0], full_arr[:,1], full_arr[:,2] = FUN_1

就节省内存而言,在这个步骤上,我认为这是我能做到的最好的,但我很乐意听取建议。无论如何,直到此时我都没有遇到问题,而且只需要大约2分钟。

以下是排序例程(有两个连续的排序)

for idx in range(2):
    sort_idx = numpy.argsort(full_arr[:,idx])
    full_arr = full_arr[sort_idx]
    # ...
    # <additional processing, return small (1000, 3) array of stats>

现在这个排序算法已经在慢慢工作了(需要大约10分钟)。然而,我最近开始在FUN_1中使用一个更大的、更细密的[X, F(X)]值表进行插值步骤,以返回b_arr。现在SORT真的变得很慢,尽管其他一切都保持不变。

有趣的是,在SORT速度变慢的步骤中,我甚至没有对插值后的值进行排序。下面是不同插值文件的一些片段 - 在每种情况下,较小的文件约小30%,在第二列的值方面更加均匀;较慢的文件具有更高的分辨率和更多独特的数据,因此插值结果可能更加独特,但我不确定这是否会产生任何影响...?

较大、较慢的文件:

17399307    99.4
17493652    98.8
17570460    98.2
17575180    97.6
17577127    97
17578255    96.4
17580576    95.8
17583028    95.2
17583699    94.6
17584172    94

更小、更均匀的常规文件:

1       24  
1001    24  
2001    24  
3001    24  
4001    24  
5001    24
6001    24
7001    24

我不确定是什么原因导致了这个问题,我很乐意听取任何关于在这种内存限制情况下进行排序的建议或一般性的意见!


4
为什么你不在构建大数组时对其进行排序 - 如在插值或合并小数组期间进行排序?你可以使用In-place Merge Sort,因为它使用O(1)内存。 - SDekov
3
如果你需要对一个超过内存大小的数据集进行排序,你可以参考归并排序。在早期,当主内存只有几千字节时,人们使用它来对磁带上的大型数据集进行排序。虽然我不太了解你的完整问题,但这可能会对你有所帮助。 - Bas Swinckels
@BasSwinckels,谢谢你,我会研究一下的。在我发现NumPy在处理大数据方面的高效性之前,我曾经试图想出一种类似于您链接文章所描述的方式将我的数据按顺序排序,尽管我没有完全理解它。困扰我的问题是大小没有改变,只是其中一个数据列的组成发生了变化,这在某种程度上阻碍了我的程序,因此我正在努力找出原因...干杯! - isosceleswheel
@YXD 嗨,不太确定你具体指的是什么... 你是指为我的数据集构建一个数据库并创建“键”或某种索引系统来返回排序后的数据吗?我觉得我描述问题不清楚,实际上我的主要问题是尝试弄清楚为什么相同大小的数组根据其内容更难排序,尽管从直觉上讲这是有道理的,因为可能需要进行更多的比较...但我真的不明白这些东西在内部是如何表示的,我还是个菜鸟 ;-) - isosceleswheel
@ali_m,谢谢你的编辑,我以后会更加小心地发布 :-) - isosceleswheel
显示剩余2条评论
2个回答

14

目前,每次调用np.argsort都会生成一个int64索引的(868940742, 1)数组,它本身将占用约7 GB的空间。此外,当您使用这些索引对full_arr的列进行排序时,会生成另一个(868940742, 1)数组的浮点数,因为花式索引始终返回副本而不是视图

有一个相当明显的改进方法是使用.sort()方法直接对full_arr进行就地排序。不幸的是,.sort()不允许您直接指定要按行或列排序的方式。然而,对于结构化数组,您可以指定要按字段排序。因此,您可以通过将数组作为具有三个浮点字段的结构化数组的视图来强制在其中一个三列上进行原地排序,然后按其中一个字段排序:

full_arr.view('f8, f8, f8').sort(order=['f0'], axis=0)

在这种情况下,我正在就第0个字段对应于第一列对full_arr进行原地排序。请注意,我假设有三个float64列('f8') - 如果您的dtype不同,您应该相应地更改此设置。这还要求您的数组是连续的并且以行为主的格式,即full_arr.flags.C_CONTIGUOUS == True
此方法的信用应归功于Joe Kington的答案here
尽管使用结构化数组按字段排序所需的内存较少,但与使用np.argsort生成索引数组相比,速度要慢得多,正如您在下面的评论中提到的那样(请参见this previous question)。如果您使用np.argsort获取一组要排序的索引,则可以通过使用np.take而不是直接索引来获得排序后的数组,从而获得适度的性能提升。
 %%timeit -n 1 -r 100 x = np.random.randn(10000, 2); idx = x[:, 0].argsort()
x[idx]
# 1 loops, best of 100: 148 µs per loop

 %%timeit -n 1 -r 100 x = np.random.randn(10000, 2); idx = x[:, 0].argsort()
np.take(x, idx, axis=0)
# 1 loops, best of 100: 42.9 µs per loop

然而,我不会期望在内存使用方面看到任何差异,因为这两种方法都会生成一份副本。

关于您关于为什么排序第二个数组更快的问题 - 是的,您应该期望任何合理的排序算法在数组中有较少唯一值时更快,因为平均而言它需要做的工作较少。假设我有一个介于1和10之间的随机数字序列:

5  1  4  8  10  2  6  9  7  3

这些数字有10! = 3628800种可能的排列方式,但只有一种方式是按升序排列的。现在假设只有5个不同的数字:

4  4  3  2  3  1  2  5  1  5

现在有 2⁵ = 32 种方式可以按升序排列这些数字,因为我可以交换排序向量中的任意一对相同数字而不破坏排序。

默认情况下,np.ndarray.sort() 使用快速排序。此算法的qsort变体通过递归地选择数组中的“枢轴”元素,然后重新排序数组,使所有小于枢轴值的元素都放在它之前,而所有大于枢轴值的元素都放在它之后。等于枢轴的值已经排好序了。具有较少唯一值意味着平均而言,在任何给定的扫描中,更多的值将等于枢轴值,因此需要更少的扫描来完全排序数组。

例如:

%%timeit -n 1 -r 100 x = np.random.random_integers(0, 10, 100000)
x.sort()
# 1 loops, best of 100: 2.3 ms per loop

%%timeit -n 1 -r 100 x = np.random.random_integers(0, 1000, 100000)
x.sort()
# 1 loops, best of 100: 4.62 ms per loop

在这个例子中,两个数组的数据类型是相同的。如果您较小的数组与较大的数组相比具有较小的项目大小,则由于使用花式索引而复制它的成本也将更小。

我认为我在问题的最后一段应该再详细解释一下,那些细节旨在帮助识别问题所在,但我并不是在寻找最佳解决方案; 而是要了解问题的性质。 我感兴趣的是,为什么对两个大小和格式完全相同的数组进行排序会导致如此不同的内存需求。 唯一的区别是“较慢”的数组具有更大的唯一值集。 最终,似乎sort([2,1,1,4,3])sort([1.66453, 2.899, 4.9034, 1.43235, 2.3464])需要更多资源,我很好奇其中的原因... :-) - isosceleswheel
顺便说一下(空间不够了!),感谢Joe的建议,让我尝试full_arr.view('f8, f8, f8').sort(order=['f0'], axis=0),我以前没有见过这种使用view的方法。 - isosceleswheel
不幸的是,结构化数组方法比其他更占用内存的方法慢得多,即使只使用 full_arr 的前10%,我也会超时。我记得在这个网站的其他地方看到过,结构化数组在各种工作中的缺点是相当慢的。 - isosceleswheel

8

编辑:如果有编程新手和使用numpy的人看到此帖子,我想指出考虑使用的np.dtype的重要性。在我的情况下,我实际上可以使用半精度浮点数,即np.float16,从而将一个20GB内存对象减少到5GB,并使排序更加可管理。由numpy使用的默认值是np.float64,这是很多您可能不需要的精度。在这里查看文档,其中描述了不同数据类型的容量。感谢@ali_m在评论中指出这一点。

我没有好好解释这个问题,但我已经发现了一些有用的解决方法,我认为对于任何需要对大型numpy数组进行排序的人都会有所帮助。

我正在从包含元素[position, value]的22个人类基因组数据“子数组”构建非常大的numpy数组。最终,必须根据特定列中的值对最终数组进行数字排序,并且不能在行内洗牌值。

子阵列的尺寸如下:

arr1.shape = (N1, 2)
...
arr22.shape = (N22, 2)

sum([N1..N2]) = 868940742,即需要排序的位置接近10亿。

首先我使用函数process_sub_arrs处理22个子数组,该函数返回一个三元组,其三个一维数组与输入长度相同。我将这些一维数组堆叠成一个新的(N, 3)数组,并将它们插入到初始化为整个数据集的np.zeros数组中:

    full_arr = np.zeros([868940742, 3])
    i, j = 0, 0

    for arr in list(arr1..arr22):  
        # indices (i, j) incremented at each loop based on sub-array size
        j += len(arr)
        full_arr[i:j, :] = np.column_stack( process_sub_arrs(arr) )
        i = j

    return full_arr

编辑:由于我意识到我的数据集可以使用半精度浮点数表示,现在我将full_arr初始化如下:full_arr = np.zeros([868940742, 3], dtype=np.float16),这只有原来大小的四分之一,更容易排序。

结果是一个巨大的20GB数组:

full_arr.nbytes = 20854577808

正如@ali_m在他的详细帖子中指出的那样,我的早期例程效率低下:

sort_idx = np.argsort(full_arr[:,idx])
full_arr = full_arr[sort_idx]

数组 sort_idx 的大小为 full_arr 的33%,在对 full_arr 进行排序后会浪费内存。由于“花哨”的索引,这种排序可能会生成 full_arr 的副本,潜在地将内存使用推到已用于容纳大数组的233%!这是一个缓慢的步骤,持续约十分钟,并且严重依赖虚拟内存。
然而,我不确定“花哨”的排序是否会生成持久副本。观察我的机器上的内存使用情况,似乎 full_arr = full_arr[sort_idx] 删除了对未排序原始数组的引用,因为大约1秒后,只剩下排序后的数组和索引所使用的内存,即使存在短暂的副本。
更紧凑的使用 argsort() 以节省内存的方法如下:
    full_arr = full_arr[full_arr[:,idx].argsort()]

这仍会在分配时产生一个峰值,其中会创建一个瞬态指数数组和一个瞬态副本,但内存几乎立即被释放。
@ali_m指出了一个很好的技巧(由Joe Kington创造),用full_arr上的view生成事实上的结构化数组。好处是这些可以“原地”排序,保持稳定的行顺序。
full_arr.view('f8, f8, f8').sort(order=['f0'], axis=0)

视图在执行数学数组操作方面表现出色,但对于排序来说,即使是从我的数据集中仅涉及一个子数组,它也远远不够高效。总的来说,结构化数组看起来并不很擅长扩展,尽管它们具有非常有用的属性。如果有人知道原因,请告诉我。

一个好的选择是最小化内存消耗并通过构建一系列小型简单函数来提高大型数组的性能。函数在完成后清除本地变量,因此如果中间数据结构正在增加并消耗内存,则这可能是一个不错的解决方案。

这是我用来加速大规模数组排序的简要流程:

def process_sub_arrs(arr):
    """process a sub-array and return a 3-tuple of 1D values arrays"""

    return values1, values2, values3

def build_arr():
    """build the initial array by joining processed sub-arrays"""

    full_arr = np.zeros([868940742, 3])
    i, j = 0, 0

    for arr in list(arr1..arr22):  
        # indices (i, j) incremented at each loop based on sub-array size
        j += len(arr)
        full_arr[i:j, :] = np.column_stack( process_sub_arrs(arr) )
        i = j

    return full_arr

def sort_arr():
    """return full_arr and sort_idx"""

    full_arr = build_arr()
    sort_idx = np.argsort(full_arr[:, index])

    return full_arr[sort_idx]

def get_sorted_arr():
    """call through nested functions to return the sorted array"""

    sorted_arr = sort_arr()
    <process sorted_arr>

    return statistics

调用栈:get_sorted_arr --> sort_arr --> build_arr --> process_sub_arrs

每个内部函数完成后,get_sorted_arr() 最终只保存排好序的数组,然后返回一小组统计数据。

编辑:在这里值得指出的是,即使您能够使用更紧凑的 dtype 来表示您的大型数组,您仍将希望在摘要计算中使用更高的精度。例如,由于 full_arr.dtype = np.float16,命令 np.mean(full_arr[:,idx]) 尝试以半精度浮点数计算平均值,但当对大量的数组求和时,很快就会溢出。使用 np.mean(full_arr[:,idx], dtype=np.float64) 可以防止溢出。

我最初发布这个问题是因为我困惑于一个相同大小的数据集突然占用了系统内存,尽管新的“慢速”集合中唯一值的比例有很大的差异。@ali_m 指出,确实如此,具有较少唯一值的更均匀的数据更容易进行排序:

Quicksort 的 qsort 变体通过递归地选择数组中的 'pivot' 元素,然后重新排列数组,以便将所有小于 pivot 值的元素放在其前面,将所有大于 pivot 值的元素放在其后面。等于 pivot 的值已经排序,因此直观地说,数组中唯一值越少,需要进行交换的次数就越少。

在这方面,我最终做出的改变是提前舍入新数据集,因为从插值步骤中留下了不必要的高十进制精度。这最终比其他内存节省步骤产生了更大的影响,显示排序算法本身是这种情况的限制因素。

期待其他人对这个主题的评论或建议,并且我几乎肯定会在某些技术问题上失言,所以我很高兴收到回音 :-)


2
如果您愿意对数据进行四舍五入,那么我建议考虑切换到一个具有较小项大小的dtype,以减少内存需求。如果您仍然需要存储小数值,则可以切换到不太精确的浮点数,例如np.float64 --> np.float32。如果您只需要整数值,则可以切换到较小的整数类型,例如np.int32/np.uint32,甚至是np.int16/np.uint16,这取决于您需要表示的最大值。如果您还没有尝试过,请尝试使用np.take而不是数组索引来获取排序后的数组。 - ali_m
嗨@ali_m,再次感谢您提供的所有有用反馈。您知道我在哪里可以找到有关浮点数和整数的dtypes限制的清晰讨论吗?我正在谷歌上搜索,并找到了有关float64 vs. float32特定应用的讨论,但我只是想了解这些类型能够表示的范围,例如,是float64仅到小数点后16位,对于float32而言则只有8位,我假设对于16位则为4位?int32 vs. int64实施的限制是什么?实际上,full_arr的3列精度相当低,最多只有其中一列达到6位小数。 - isosceleswheel
@ali_m 其实文档里就有:http://docs.scipy.org/doc/numpy/user/basics.types.html - isosceleswheel

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接