numpy argsort 性能下降

Question

如果这个问题看起来很基本，我真诚地道歉。

鉴于：

>>> import numpy as np
>>> import time
>>> A = np.random.rand( int(1e5), int(5e4) ) # large numpy array

目标：

>>> bt=time.time(); B=np.argsort(A,axis=1);et=time.time();print(f"Took {(et-bt):.2f} s")

但是，计算索引数组需要很长时间：

# Took 316.90 s

题：

有没有其他时间有效的方法来做到这一点？

干杯，

Answer 1

输入数组A的形状为(100_000, 50_000)并且默认包含np.float64值。 这意味着您需要8 * 100_000 * 50_000 / 1024**3 = 37.2 Gio的内存仅用于此数组。 您可能还需要为输出矩阵B （它应该包含类型为np.int64项目）提供相同数量的空间。 这意味着您需要一台至少具有 74.4 Gio 的机器，更不用说操作系统 (OS) 和运行软件所需的空间（因此可能至少为 80 Gio）。 如果您没有这样的内存空间，那么您的操作系统将使用您的存储设备作为交换内存，这会慢得多。

假设你有这样一个可用的内存空间，这样的计算是非常昂贵的。 这主要是由于填充B数组时出现页面错误，以及B非常大以及 Numpy 的默认实现按顺序进行计算的事实。 您可以使用并行 Numba 代码和较小的输出来加快计算速度。 下面是一个例子：

import numba as nb

@nb.njit('int32[:,::1](float64[:,::1])', parallel=True)
def fastSort(a):
    b = np.empty(a.shape, dtype=np.int32)
    for i in nb.prange(a.shape[0]):
        b[i,:] = np.argsort(a[i,:])
    return b

请注意，Numba的实施argsort比与NumPy的之一，但其水货版本低效率要快很多，如果目标机器有MNY核心和良好的内存带宽。

以下是我的 6 核机器在大小矩阵(10_000, 50_000) （小 10 倍，因为我没有 80 Gio 的 RAM）：

Original implementation: 28.14 s
Sequential Numba:        38.70 s
Parallel Numba:           6.79 s

因此，所得解决方案的速度提高了 4.1 倍。

请注意，在这种特定情况下，您甚至可以对B的项目使用类型uint16 ，因为每行的大小小于2**16 = 65536 。 这可能不会明显更快，但它应该节省一些额外的内存。 由此产生的所需内存将为 46.5 Gio。 您可以使用np.float32类型进一步减少所需的内存量（通常以牺牲精度为代价）。

如果您想进一步缩短执行时间，那么您需要使用 C 或 C++ 等低级语言为您的特定需求实现更快的argsort实现。 但是请注意，如果您不是使用这种语言的有经验的程序员或不熟悉低级优化，那么击败 Numpy 绝非易事。 如果您对这样的解决方案感兴趣，一个好的开始可能是实现基数排序。