使用 numpy argsort 的 MemoryError答案

【问题标题】：MemoryError using numpy argsort使用 numpy argsort 的 MemoryError
【发布时间】：2015-05-07 15:13:52
【问题描述】：

我正在尝试对一个大型 numpy 数组进行排序。它是一个 4 维“float16”数组，形状为（450、470、10、470）（994,050,000 个元素）。加载时数组大小约为 2 GB，我使用的是 64 位 linux（和 64 位 python），所以我应该能够毫无问题地使用所有 8GB 的 RAM，但是一旦我尝试使用 argsort，我得到一个内存错误。

代码如下：

import numpy as np

score = np.load('score.npy')
aaarg = np.argsort(score, axis=None)[:-21:-1]

我的目标是获取数组中最大 20 个元素的索引。我也尝试使用 argpartition 而不是 argsort 但我遇到了同样的问题。有什么办法吗？

我正在使用 python 2.7.9 (Anaconda 2.2.0) 和 numpy 1.9.2 谢谢。

【问题讨论】：

你可以用np.savetxt把它分成几部分，对部分进行排序，然后得到最高的n个值
你试过score.sort()吗？由于它进行就地排序，因此它的内存效率可能要高得多。我猜argsort() 需要的内存至少是score 的三倍（原始、要排序的副本和返回的索引）。
一次对数组的一部分进行排序，然后对结果进行排序，谢谢你的想法。因为我需要索引，所以不能只使用 sort()。

标签： python python-2.7 numpy anaconda

【解决方案1】：

首先应该指出，argsort 返回一个索引数组，将数据索引到一个排序数组中。例如：

a = np.random.random(100)
idx = a.argsort()
b = a[idx]

如果您使用 64 位 python 和 numpy，idx 的 dtype 将是 int64，所以让我们做一些数学运算。 994,050,000 * 8 字节 / 1024**3 字节 / GB = 7.4 GB。这不包括 score 数组或任何系统/python 开销。

虽然您应该能够进行就地排序，但您将无法在内存中对该大小的数组进行 argsort。

【讨论】：

啊，好吧，我希望只取索引数组的最后 20 个元素（使用 [:-21:-1] ），我会避免这个问题。
@MicheleValentini 在这种特殊情况下，您可以使用很多技巧。只有“2**16”个唯一的 float16 值，因此您永远不需要“排序”其中的 1e9 个。你可以用numpy.bincount 来代替一些花哨的技巧。如果您只需要最大的 20 个值，您还可以使用 argmax，它应该相当快并且不需要额外的内存。