python - 使用 numpy argsort 出现内存错误

Question

我正在尝试对一个大的 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谢谢。

Answer 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 开销。

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